openstat-cli 1.0.0__py3-none-any.whl

openstat/io/loader.py

@@ -0,0 +1,187 @@
+"""Data loading and saving: CSV, Parquet, Stata (.dta), Excel (.xlsx)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import polars as pl
+
+from openstat.config import get_config
+from openstat.logging_config import get_logger
+
+log = get_logger("io")
+
+
+def _load_dta(path: Path, session=None) -> pl.DataFrame:
+    """Load a Stata .dta file using pyreadstat (optional dependency)."""
+    try:
+        import pyreadstat
+    except ImportError:
+        raise ImportError(
+            "Reading .dta files requires pyreadstat. "
+            "Install it with: pip install openstat[stata]"
+        )
+    pandas_df, meta = pyreadstat.read_dta(str(path))
+    if session is not None and meta.variable_value_labels:
+        session._variable_labels = meta.variable_value_labels
+    return pl.from_pandas(pandas_df)
+
+
+def _load_sas7bdat(path: Path, session=None) -> pl.DataFrame:
+    """Load a SAS .sas7bdat file using pyreadstat (optional dependency)."""
+    try:
+        import pyreadstat
+    except ImportError:
+        raise ImportError(
+            "Reading .sas7bdat files requires pyreadstat. "
+            "Install it with: pip install openstat[sas]"
+        )
+    pandas_df, meta = pyreadstat.read_sas7bdat(str(path))
+    if session is not None and meta.variable_value_labels:
+        session._variable_labels = meta.variable_value_labels
+    return pl.from_pandas(pandas_df)
+
+
+def _load_sav(path: Path, session=None) -> pl.DataFrame:
+    """Load an SPSS .sav file using pyreadstat (optional dependency)."""
+    try:
+        import pyreadstat
+    except ImportError:
+        raise ImportError(
+            "Reading .sav files requires pyreadstat. "
+            "Install it with: pip install openstat[spss]"
+        )
+    pandas_df, meta = pyreadstat.read_sav(str(path))
+    if session is not None and meta.variable_value_labels:
+        session._variable_labels = meta.variable_value_labels
+    return pl.from_pandas(pandas_df)
+
+
+def _load_excel(path: Path, sheet: str | int | None = None) -> pl.DataFrame:
+    """Load an Excel file (optional dependency).
+
+    Args:
+        sheet: Sheet name or 0-based index. If None, loads the first sheet.
+               Pass 'list' to list available sheet names.
+    """
+    try:
+        import openpyxl
+    except ImportError:
+        raise ImportError(
+            "Reading .xlsx files requires openpyxl. "
+            "Install it with: pip install openstat[excel]"
+        )
+
+    if sheet == "list":
+        wb = openpyxl.load_workbook(str(path), read_only=True, data_only=True)
+        names = wb.sheetnames
+        wb.close()
+        raise ValueError(f"Sheets in {path.name}: {', '.join(names)}")
+
+    kwargs: dict = {}
+    if sheet is not None:
+        kwargs["sheet_name"] = sheet
+    try:
+        return pl.read_excel(path, **kwargs)
+    except TypeError:
+        # Older polars may not support sheet_name
+        return pl.read_excel(path)
+
+
+def _load_csv(path: Path, session=None) -> pl.DataFrame:
+    cfg = get_config()
+    return pl.read_csv(
+        path,
+        infer_schema_length=cfg.infer_schema_length,
+        separator=cfg.csv_separator,
+    )
+
+
+_LOADERS = {
+    ".csv": _load_csv,
+    ".parquet": lambda p, session=None: pl.read_parquet(p),
+    ".dta": _load_dta,
+    ".xlsx": lambda p, session=None: _load_excel(p),
+    ".xls": lambda p, session=None: _load_excel(p),
+    ".sas7bdat": _load_sas7bdat,
+    ".sav": _load_sav,
+}
+
+_SUPPORTED = ", ".join(_LOADERS.keys())
+
+
+def load_file(path: str | Path, session=None) -> pl.DataFrame:
+    """Load a data file into a Polars DataFrame."""
+    p = Path(path)
+    if not p.exists():
+        raise FileNotFoundError(f"File not found: {p}")
+
+    suffix = p.suffix.lower()
+    loader = _LOADERS.get(suffix)
+    if loader is None:
+        raise ValueError(f"Unsupported file format: {suffix}  (supported: {_SUPPORTED})")
+    log.info("Loading %s (%s, %.1f KB)", p.name, suffix, p.stat().st_size / 1024)
+    df = loader(p, session=session)
+    log.info("Loaded %d rows x %d cols", df.height, df.width)
+    return df
+
+
+def _save_excel(df: pl.DataFrame, path: Path) -> None:
+    """Save to Excel (optional dependency)."""
+    try:
+        df.write_excel(path)
+    except (ImportError, ModuleNotFoundError):
+        raise ImportError(
+            "Writing .xlsx files requires xlsxwriter. "
+            "Install it with: pip install openstat[excel]"
+        )
+
+
+def _save_dta(df: pl.DataFrame, path: Path) -> None:
+    """Save to Stata .dta (optional dependency)."""
+    try:
+        import pyreadstat
+    except ImportError:
+        raise ImportError(
+            "Writing .dta files requires pyreadstat. "
+            "Install it with: pip install openstat[stata]"
+        )
+    pandas_df = df.to_pandas()
+    pyreadstat.write_dta(pandas_df, str(path))
+
+
+def _save_sav(df: pl.DataFrame, path: Path) -> None:
+    """Save to SPSS .sav (optional dependency)."""
+    try:
+        import pyreadstat
+    except ImportError:
+        raise ImportError(
+            "Writing .sav files requires pyreadstat. "
+            "Install it with: pip install openstat[spss]"
+        )
+    pandas_df = df.to_pandas()
+    pyreadstat.write_sav(pandas_df, str(path))
+
+
+def save_file(df: pl.DataFrame, path: str | Path) -> Path:
+    """Save a DataFrame to CSV, Parquet, Excel, Stata, or SPSS."""
+    p = Path(path)
+    p.parent.mkdir(parents=True, exist_ok=True)
+
+    suffix = p.suffix.lower()
+    if suffix == ".csv":
+        df.write_csv(p, separator=get_config().csv_separator)
+    elif suffix == ".parquet":
+        df.write_parquet(p)
+    elif suffix == ".xlsx":
+        _save_excel(df, p)
+    elif suffix == ".dta":
+        _save_dta(df, p)
+    elif suffix == ".sav":
+        _save_sav(df, p)
+    else:
+        raise ValueError(
+            f"Unsupported save format: {suffix}  "
+            "(use .csv, .parquet, .xlsx, .dta, or .sav)"
+        )
+    return p

openstat/jupyter/__init__.py

@@ -0,0 +1,18 @@
+"""Jupyter notebook integration for OpenStat.
+
+Usage in Jupyter:
+    %load_ext openstat
+    %ost load data.csv
+    %ost summarize
+
+    %%openstat
+    load data.csv
+    summarize
+    ols y ~ x1 + x2
+"""
+
+
+def load_ipython_extension(ipython):
+    """Register OpenStat magics when %load_ext openstat is called."""
+    from openstat.jupyter.magic import OpenStatMagics
+    ipython.register_magics(OpenStatMagics)

openstat/jupyter/display.py

@@ -0,0 +1,18 @@
+"""Jupyter display helpers for OpenStat."""
+
+from __future__ import annotations
+
+
+def fit_result_to_html(result) -> str:
+    """Convert a FitResult to HTML table for Jupyter rendering."""
+    from rich.console import Console
+    console = Console(record=True, width=120)
+    text = result.summary_table()
+    console.print(text)
+    return console.export_html(inline_styles=True)
+
+
+def dataframe_to_html(df, max_rows: int = 50) -> str:
+    """Convert a Polars DataFrame to styled HTML for Jupyter."""
+    pdf = df.head(max_rows).to_pandas()
+    return pdf.to_html(classes="openstat-table", border=0, index=False)

openstat/jupyter/magic.py

@@ -0,0 +1,60 @@
+"""IPython magic commands for OpenStat in Jupyter notebooks."""
+
+from __future__ import annotations
+
+try:
+    from IPython.core.magic import Magics, magics_class, cell_magic, line_magic
+    from IPython.display import display, HTML
+    HAS_IPYTHON = True
+except ImportError:
+    HAS_IPYTHON = False
+
+from openstat.session import Session
+
+
+def _dispatch_line(session: Session, line: str) -> str | None:
+    """Dispatch a single command line using the REPL dispatcher."""
+    from openstat.repl import _dispatch
+    return _dispatch(session, line)
+
+
+def _rich_to_html(text: str) -> str:
+    """Convert Rich-styled text to HTML."""
+    from rich.console import Console
+    console = Console(record=True, width=120)
+    console.print(text)
+    return console.export_html(inline_styles=True)
+
+
+if HAS_IPYTHON:
+    @magics_class
+    class OpenStatMagics(Magics):
+        """IPython magics for OpenStat: %ost and %%openstat."""
+
+        def __init__(self, shell):
+            super().__init__(shell)
+            self.session = Session()
+
+        @line_magic
+        def ost(self, line):
+            """Run a single OpenStat command: %ost load data.csv"""
+            result = _dispatch_line(self.session, line)
+            if result and result != "__QUIT__":
+                display(HTML(f"<pre>{_rich_to_html(result)}</pre>"))
+
+        @cell_magic
+        def openstat(self, line, cell):
+            """Run multiple OpenStat commands in a cell."""
+            for cmd_line in cell.strip().split('\n'):
+                cmd_line = cmd_line.strip()
+                if not cmd_line or cmd_line.startswith('#'):
+                    continue
+                result = _dispatch_line(self.session, cmd_line)
+                if result == "__QUIT__":
+                    break
+                if result:
+                    display(HTML(f"<pre>{_rich_to_html(result)}</pre>"))
+else:
+    class OpenStatMagics:
+        """Stub when IPython is not available."""
+        pass

openstat/logging_config.py

@@ -0,0 +1,59 @@
+"""Logging configuration for OpenStat."""
+
+from __future__ import annotations
+
+import logging
+import sys
+from pathlib import Path
+
+_LOG_DIR = Path.home() / ".openstat" / "logs"
+_LOG_FORMAT = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+_LOG_DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+_configured = False
+
+
+def setup_logging(*, verbose: bool = False, debug: bool = False) -> None:
+    """Configure logging for OpenStat.
+
+    - Default: WARNING to file only
+    - --verbose: INFO to file + stderr
+    - --debug: DEBUG to file + stderr
+    """
+    global _configured
+    if _configured:
+        return
+    _configured = True
+
+    root = logging.getLogger("openstat")
+
+    if debug:
+        level = logging.DEBUG
+    elif verbose:
+        level = logging.INFO
+    else:
+        level = logging.WARNING
+
+    root.setLevel(logging.DEBUG)  # capture everything, handlers filter
+
+    # File handler — always active, always DEBUG
+    try:
+        _LOG_DIR.mkdir(parents=True, exist_ok=True)
+        fh = logging.FileHandler(_LOG_DIR / "openstat.log", encoding="utf-8")
+        fh.setLevel(logging.DEBUG)
+        fh.setFormatter(logging.Formatter(_LOG_FORMAT, datefmt=_LOG_DATE_FORMAT))
+        root.addHandler(fh)
+    except OSError:
+        pass  # if we can't write logs, don't crash
+
+    # Console handler — only if verbose/debug
+    if verbose or debug:
+        ch = logging.StreamHandler(sys.stderr)
+        ch.setLevel(level)
+        ch.setFormatter(logging.Formatter(_LOG_FORMAT, datefmt=_LOG_DATE_FORMAT))
+        root.addHandler(ch)
+
+
+def get_logger(name: str) -> logging.Logger:
+    """Get a logger under the openstat namespace."""
+    return logging.getLogger(f"openstat.{name}")