openstat-cli 1.0.0__py3-none-any.whl

openstat/commands/stata_import_cmds.py

@@ -0,0 +1,215 @@
+"""Stata .do file importer: convert Stata syntax to OpenStat .ost script."""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+from openstat.commands.base import command, CommandArgs, friendly_error
+from openstat.session import Session
+
+# ── Stata → OpenStat translation rules ──────────────────────────────────────
+# Each rule is (pattern, replacement_or_callable)
+# Applied in order; first match wins for each line.
+
+_RULES: list[tuple[re.Pattern, object]] = []
+
+
+def _rule(pattern: str, repl):
+    _RULES.append((re.compile(pattern, re.IGNORECASE), repl))
+
+
+# Comments
+_rule(r"^\s*\*.*$", lambda m: "# " + m.group(0).lstrip("* \t"))
+_rule(r"^(\s*)//(.*)$", lambda m: m.group(1) + "# " + m.group(2))
+
+# use → load
+_rule(r"^\s*use\s+[\"']?(.+?)[\"']?\s*(?:,\s*clear)?\s*$",
+      lambda m: f"load {m.group(1).strip()}")
+
+# save → save
+_rule(r"^\s*save\s+[\"']?(.+?)[\"']?\s*(?:,\s*replace)?\s*$",
+      lambda m: f"save {m.group(1).strip()}")
+
+# summarize → summarize
+_rule(r"^\s*su(?:mmarize)?\s*(.*)", lambda m: f"summarize {m.group(1).strip()}")
+
+# describe → describe
+_rule(r"^\s*desc(?:ribe)?\s*(.*)", lambda m: f"describe")
+
+# regress → ols
+_rule(r"^\s*reg(?:ress)?\s+(\S+)\s+(.*)",
+      lambda m: f"ols {m.group(1)} {m.group(2).strip()}")
+
+# logit → logit (same)
+_rule(r"^\s*logit\s+(\S+)\s+(.*)",
+      lambda m: f"logit {m.group(1)} {m.group(2).strip()}")
+
+# probit → probit (same)
+_rule(r"^\s*probit\s+(\S+)\s+(.*)",
+      lambda m: f"probit {m.group(1)} {m.group(2).strip()}")
+
+# poisson → poisson (same)
+_rule(r"^\s*poisson\s+(\S+)\s+(.*)",
+      lambda m: f"poisson {m.group(1)} {m.group(2).strip()}")
+
+# tab → tabulate
+_rule(r"^\s*tab(?:ulate)?\s+(.*)", lambda m: f"tabulate {m.group(1).strip()}")
+
+# cor → correlate
+_rule(r"^\s*cor(?:relate)?\s+(.*)", lambda m: f"correlate {m.group(1).strip()}")
+
+# drop if → filter
+_rule(r"^\s*drop\s+if\s+(.*)",
+      lambda m: f"# [manual] drop if {m.group(1)}  →  filter {_stata_cond(m.group(1))}")
+
+# keep if → filter
+_rule(r"^\s*keep\s+if\s+(.*)",
+      lambda m: f"# [manual] keep if {m.group(1)}  →  filter {_stata_cond(m.group(1))}")
+
+# drop varlist → drop
+_rule(r"^\s*drop\s+(?!if)(.*)", lambda m: f"drop {m.group(1).strip()}")
+
+# keep varlist → keep
+_rule(r"^\s*keep\s+(?!if)(.*)", lambda m: f"keep {m.group(1).strip()}")
+
+# gen/generate → generate
+_rule(r"^\s*gen(?:erate)?\s+(.+)=(.+)",
+      lambda m: f"generate {m.group(1).strip()} = {m.group(2).strip()}")
+
+# replace → replace
+_rule(r"^\s*replace\s+(.+)=(.+)",
+      lambda m: f"replace {m.group(1).strip()} = {m.group(2).strip()}")
+
+# rename → rename
+_rule(r"^\s*rename\s+(\S+)\s+(\S+)", lambda m: f"rename {m.group(1)} {m.group(2)}")
+
+# sort → sort
+_rule(r"^\s*sort\s+(.*)", lambda m: f"sort {m.group(1).strip()}")
+
+# set seed → set seed
+_rule(r"^\s*set\s+seed\s+(\d+)", lambda m: f"set seed {m.group(1)}")
+
+# xtset → xtset (panel)
+_rule(r"^\s*xtset\s+(.*)", lambda m: f"xtset {m.group(1).strip()}")
+
+# tsset → tsset
+_rule(r"^\s*tsset\s+(.*)", lambda m: f"tsset {m.group(1).strip()}")
+
+# stset → stset
+_rule(r"^\s*stset\s+(.*)", lambda m: f"stset {m.group(1).strip()}")
+
+# display → print (closest)
+_rule(r"^\s*dis(?:play)?\s+(.*)", lambda m: f"# display {m.group(1).strip()}")
+
+# local macro → define
+_rule(r"^\s*local\s+(\w+)\s+(.*)", lambda m: f"define {m.group(1)} = {m.group(2).strip()}")
+
+# forvalues
+_rule(r"^\s*forvalues\s+(.*)", lambda m: f"forvalues {m.group(1).strip()} {{")
+
+# foreach
+_rule(r"^\s*foreach\s+(.*)", lambda m: f"foreach {m.group(1).strip()} {{")
+
+# if / else blocks - pass through
+_rule(r"^\s*\}\s*$", lambda m: "}")
+
+# quietly → strip (run silently)
+_rule(r"^\s*quietly\s+(.*)", lambda m: m.group(1).strip())
+
+# capture → strip (ignore errors)
+_rule(r"^\s*capture\s+(.*)", lambda m: m.group(1).strip())
+
+
+def _stata_cond(cond: str) -> str:
+    """Very rough Stata condition → Polars filter hint."""
+    # == → ==, != → !=, & → &, | → |
+    cond = cond.replace(" & ", " & ").replace(" | ", " | ")
+    return cond
+
+
+def _translate_line(line: str) -> str:
+    """Translate a single Stata .do line to OpenStat syntax."""
+    stripped = line.rstrip()
+    if not stripped.strip():
+        return ""
+
+    for pat, repl in _RULES:
+        m = pat.match(stripped)
+        if m:
+            if callable(repl):
+                return repl(m)
+            return repl
+    # No match — keep as comment with note
+    return f"# [untranslated] {stripped}"
+
+
+def convert_do_file(do_text: str) -> str:
+    """Convert Stata .do file content to OpenStat .ost script."""
+    out_lines = ["# OpenStat script — converted from Stata .do file", ""]
+    for line in do_text.splitlines():
+        translated = _translate_line(line)
+        out_lines.append(translated)
+    return "\n".join(out_lines)
+
+
+@command("import do", usage="import do <file.do> [--out=<file.ost>]")
+def cmd_import_do(session: Session, args: str) -> str:
+    """Convert a Stata .do file to an OpenStat .ost script.
+
+    Translates common Stata commands (use, regress, gen, etc.) to their
+    OpenStat equivalents. Untranslated lines are kept as comments.
+
+    Options:
+      --out=<path>   output .ost file path (default: same name, .ost extension)
+      --run          also execute the converted script immediately
+
+    Examples:
+      import do analysis.do
+      import do stata_script.do --out=my_analysis.ost --run
+    """
+    ca = CommandArgs(args)
+    if not ca.positional:
+        return "Usage: import do <file.do> [--out=<path>] [--run]"
+
+    do_path = Path(ca.positional[0])
+    if not do_path.exists():
+        return f"File not found: {do_path}"
+
+    out_path_str = ca.options.get("out", str(do_path.with_suffix(".ost")))
+    out_path = Path(out_path_str)
+    run_after = "--run" in args
+
+    try:
+        do_text = do_path.read_text(encoding="utf-8", errors="replace")
+        ost_text = convert_do_file(do_text)
+
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        out_path.write_text(ost_text, encoding="utf-8")
+
+        n_lines = len(do_text.splitlines())
+        n_translated = sum(
+            1 for line in ost_text.splitlines()
+            if line.strip() and not line.startswith("# [untranslated]")
+            and not line.startswith("# OpenStat")
+        )
+        n_untranslated = ost_text.count("# [untranslated]")
+
+        result = (
+            f"Converted: {do_path} → {out_path}\n"
+            f"  Lines: {n_lines}  |  Translated: {n_translated}  |  "
+            f"Untranslated (kept as comments): {n_untranslated}"
+        )
+
+        if run_after:
+            from openstat.script_runner import run_script_advanced
+            try:
+                run_script_advanced(str(out_path), session)
+                result += f"\nScript executed: {out_path}"
+            except Exception as exc:
+                result += f"\nScript error: {exc}"
+
+        return result
+
+    except Exception as e:
+        return friendly_error(e, "import do")

openstat/commands/string_cmds.py

@@ -0,0 +1,124 @@
+"""String and date manipulation commands."""
+
+from __future__ import annotations
+
+import re
+
+import polars as pl
+
+from openstat.commands.base import command
+from openstat.session import Session
+
+
+def _stata_opts(raw: str) -> tuple[list[str], dict[str, str]]:
+    opts: dict[str, str] = {}
+    for m in re.finditer(r'(\w+)\(([^)]*)\)', raw):
+        opts[m.group(1).lower()] = m.group(2)
+    rest = re.sub(r'\w+\([^)]*\)', '', raw)
+    positional = [t.strip(',') for t in rest.split() if t.strip(',')]
+    return positional, opts
+
+
+@command("split", usage="split varname [, sep(,) gen(newvar)]")
+def cmd_split(session: Session, args: str) -> str:
+    """Split a string column into multiple columns."""
+    df = session.require_data()
+    positional, opts = _stata_opts(args)
+    if not positional:
+        return "Usage: split varname [sep(,) gen(prefix)]"
+    var = positional[0]
+    if var not in df.columns:
+        return f"Column '{var}' not found."
+    sep = opts.get("sep", " ")
+    prefix = opts.get("gen", var)
+    session.snapshot()
+    try:
+        parts = df[var].str.split(sep)
+        max_parts = max(len(p) for p in parts.to_list())
+        new_df = df
+        for i in range(max_parts):
+            col_name = f"{prefix}{i+1}"
+            new_df = new_df.with_columns(
+                parts.list.get(i, null_on_oob=True).alias(col_name)
+            )
+        session.df = new_df
+        return f"Split '{var}' into {max_parts} columns: {[f'{prefix}{i+1}' for i in range(max_parts)]}"
+    except Exception as exc:
+        return f"split error: {exc}"
+
+
+@command("strtrim", usage="strtrim varname [gen(newvar)]")
+def cmd_strtrim(session: Session, args: str) -> str:
+    """Trim whitespace from a string column."""
+    df = session.require_data()
+    positional, opts = _stata_opts(args)
+    if not positional:
+        return "Usage: strtrim varname [gen(newvar)]"
+    var = positional[0]
+    if var not in df.columns:
+        return f"Column '{var}' not found."
+    new_var = opts.get("gen", var)
+    session.snapshot()
+    try:
+        session.df = df.with_columns(pl.col(var).str.strip_chars().alias(new_var))
+        return f"Trimmed '{var}' → '{new_var}'"
+    except Exception as exc:
+        return f"strtrim error: {exc}"
+
+
+@command("strupper", usage="strupper varname [gen(newvar)]")
+def cmd_strupper(session: Session, args: str) -> str:
+    """Convert string column to uppercase."""
+    df = session.require_data()
+    positional, opts = _stata_opts(args)
+    if not positional:
+        return "Usage: strupper varname [gen(newvar)]"
+    var = positional[0]
+    if var not in df.columns:
+        return f"Column '{var}' not found."
+    new_var = opts.get("gen", var)
+    session.snapshot()
+    try:
+        session.df = df.with_columns(pl.col(var).str.to_uppercase().alias(new_var))
+        return f"Uppercased '{var}' → '{new_var}'"
+    except Exception as exc:
+        return f"strupper error: {exc}"
+
+
+@command("strlower", usage="strlower varname [gen(newvar)]")
+def cmd_strlower(session: Session, args: str) -> str:
+    """Convert string column to lowercase."""
+    df = session.require_data()
+    positional, opts = _stata_opts(args)
+    if not positional:
+        return "Usage: strlower varname [gen(newvar)]"
+    var = positional[0]
+    if var not in df.columns:
+        return f"Column '{var}' not found."
+    new_var = opts.get("gen", var)
+    session.snapshot()
+    try:
+        session.df = df.with_columns(pl.col(var).str.to_lowercase().alias(new_var))
+        return f"Lowercased '{var}' → '{new_var}'"
+    except Exception as exc:
+        return f"strlower error: {exc}"
+
+
+@command("strreplace", usage="strreplace varname old new [gen(newvar)]")
+def cmd_strreplace(session: Session, args: str) -> str:
+    """Replace substring in a string column."""
+    df = session.require_data()
+    positional, opts = _stata_opts(args)
+    if len(positional) < 3:
+        return "Usage: strreplace varname old new [gen(newvar)]"
+    var, old_str, new_str = positional[0], positional[1], positional[2]
+    if var not in df.columns:
+        return f"Column '{var}' not found."
+    new_var = opts.get("gen", var)
+    session.snapshot()
+    try:
+        session.df = df.with_columns(pl.col(var).str.replace_all(old_str, new_str).alias(new_var))
+        return f"Replaced '{old_str}' with '{new_str}' in '{var}' → '{new_var}'"
+    except Exception as exc:
+        return f"strreplace error: {exc}"
+

openstat/commands/surv_cmds.py

@@ -0,0 +1,145 @@
+"""Survival analysis commands: stset, stcox, sts."""
+
+from __future__ import annotations
+
+import re
+
+from openstat.session import Session, ModelResult
+from openstat.commands.base import command, CommandArgs, friendly_error
+
+
+@command("stset", usage="stset <time_var>, failure(<event_var>)")
+def cmd_stset(session: Session, args: str) -> str:
+    """Declare survival time and failure event variables."""
+    df = session.require_data()
+
+    m = re.search(r'failure\((\w+)\)', args)
+    if not m:
+        return "Usage: stset <time_var>, failure(<event_var>)"
+
+    event_var = m.group(1)
+    # Time var is everything before the comma
+    time_part = args[:args.index(',')] if ',' in args else args[:m.start()]
+    time_var = time_part.strip()
+
+    if time_var not in df.columns:
+        return f"Column not found: {time_var}"
+    if event_var not in df.columns:
+        return f"Column not found: {event_var}"
+
+    session._surv_time_var = time_var
+    session._surv_event_var = event_var
+
+    n = df.height
+    events = df[event_var].sum()
+    return (
+        f"Survival time: {time_var}\n"
+        f"Failure event: {event_var}\n"
+        f"Observations: {n}, Events: {events}, Censored: {n - events}"
+    )
+
+
+@command("stcox", usage="stcox x1 x2 [--robust]")
+def cmd_stcox(session: Session, args: str) -> str:
+    """Fit a Cox Proportional Hazards model."""
+    df = session.require_data()
+
+    if session._surv_time_var is None or session._surv_event_var is None:
+        return "Survival structure not set. Use: stset <time_var>, failure(<event_var>)"
+
+    ca = CommandArgs(args)
+    robust = ca.has_flag("--robust")
+    covariates = [p for p in ca.positional if not p.startswith("--")]
+
+    if not covariates:
+        return "Usage: stcox x1 x2 [--robust]"
+
+    missing = [c for c in covariates if c not in df.columns]
+    if missing:
+        return f"Columns not found: {', '.join(missing)}"
+
+    try:
+        from openstat.stats.survival import fit_cox_ph
+
+        result, raw = fit_cox_ph(
+            df, session._surv_time_var, session._surv_event_var,
+            covariates, robust=robust,
+        )
+
+        session._last_model = raw
+        session._last_model_vars = (session._surv_time_var, covariates)
+        session._last_fit_result = result
+        session._last_fit_kwargs = {"survival": True}
+
+        md = result.to_markdown() if hasattr(result, "to_markdown") else ""
+        session.results.append(ModelResult(
+            name="Cox PH", formula=result.formula,
+            table=md, details={
+                "n_obs": result.n_obs,
+                "params": dict(result.params),
+                "log_likelihood": result.log_likelihood,
+            },
+        ))
+
+        output = result.summary_table()
+        if result.warnings:
+            output += "\n" + "\n".join(result.warnings)
+        return output
+    except ImportError as e:
+        return str(e)
+    except Exception as e:
+        return friendly_error(e, "stcox")
+
+
+@command("sts", usage="sts graph [by=group] | sts test <group_var>")
+def cmd_sts(session: Session, args: str) -> str:
+    """Kaplan-Meier survival curves and log-rank test."""
+    df = session.require_data()
+
+    if session._surv_time_var is None or session._surv_event_var is None:
+        return "Survival structure not set. Use: stset <time_var>, failure(<event_var>)"
+
+    ca = CommandArgs(args)
+    subcmd = ca.positional[0].lower() if ca.positional else ""
+
+    if subcmd == "graph":
+        group_var = ca.get_option("by")
+        try:
+            from openstat.stats.survival import kaplan_meier
+            summary, kmf = kaplan_meier(
+                df, session._surv_time_var, session._surv_event_var, group_var,
+            )
+
+            # Plot
+            try:
+                from openstat.plots.surv_plots import plot_km
+                path = plot_km(kmf, session.output_dir, group_var)
+                session.plot_paths.append(str(path))
+                summary += f"\nPlot saved: {path}"
+            except Exception:
+                pass
+
+            return summary
+        except ImportError as e:
+            return str(e)
+        except Exception as e:
+            return friendly_error(e, "sts graph")
+
+    elif subcmd == "test":
+        group_var = ca.positional[1] if len(ca.positional) > 1 else None
+        if not group_var:
+            return "Usage: sts test <group_var>"
+        if group_var not in df.columns:
+            return f"Column not found: {group_var}"
+        try:
+            from openstat.stats.survival import log_rank_test
+            return log_rank_test(
+                df, session._surv_time_var, session._surv_event_var, group_var,
+            )
+        except ImportError as e:
+            return str(e)
+        except Exception as e:
+            return friendly_error(e, "sts test")
+
+    else:
+        return "Usage: sts graph [by=group] | sts test <group_var>"

openstat/commands/survey_cmds.py

@@ -0,0 +1,153 @@
+"""Survey weighting commands: svyset, svy: prefix."""
+
+from __future__ import annotations
+
+import re
+
+from openstat.session import Session, ModelResult
+from openstat.commands.base import command, CommandArgs, friendly_error
+from openstat.dsl.parser import parse_formula, ParseError
+from openstat.types import NUMERIC_DTYPES
+
+
+@command("svyset", usage="svyset <psu> [pw=<weight>], strata(<strata>)")
+def cmd_svyset(session: Session, args: str) -> str:
+    """Declare survey design: PSU, sampling weights, and strata."""
+    df = session.require_data()
+
+    # Parse weight: [pw=weight_var]
+    m_pw = re.search(r'\[pw=(\w+)\]', args)
+    weight_var = m_pw.group(1) if m_pw else None
+
+    # Parse strata: strata(strata_var)
+    m_strata = re.search(r'strata\((\w+)\)', args)
+    strata_var = m_strata.group(1) if m_strata else None
+
+    # PSU is the first positional argument
+    clean = args
+    if m_pw:
+        clean = clean.replace(m_pw.group(0), "")
+    if m_strata:
+        clean = clean.replace(m_strata.group(0), "")
+    clean = clean.replace(",", "").strip()
+    psu_var = clean.split()[0] if clean.split() else None
+
+    # Validate columns exist
+    for var, label in [(psu_var, "PSU"), (weight_var, "weight"), (strata_var, "strata")]:
+        if var and var not in df.columns:
+            return f"{label} column not found: {var}"
+
+    session._svy_psu_var = psu_var
+    session._svy_weight_var = weight_var
+    session._svy_strata_var = strata_var
+
+    lines = ["Survey design set:"]
+    if psu_var:
+        lines.append(f"  PSU: {psu_var} ({df[psu_var].n_unique()} clusters)")
+    if weight_var:
+        lines.append(f"  Weight: {weight_var}")
+    if strata_var:
+        lines.append(f"  Strata: {strata_var} ({df[strata_var].n_unique()} strata)")
+    return "\n".join(lines)
+
+
+@command("svy:", usage="svy: summarize|ols|logit ...")
+def cmd_svy(session: Session, args: str) -> str:
+    """Run survey-weighted analysis. Requires svyset first."""
+    df = session.require_data()
+
+    if session._svy_weight_var is None:
+        return "Survey design not set. Use: svyset <psu> [pw=<weight>], strata(<strata>)"
+
+    parts = args.strip().split(None, 1)
+    subcmd = parts[0].lower() if parts else ""
+    rest = parts[1] if len(parts) > 1 else ""
+
+    if subcmd == "summarize":
+        return _svy_summarize(session, df, rest)
+    elif subcmd == "ols":
+        return _svy_ols(session, df, rest)
+    elif subcmd == "logit":
+        return _svy_logit(session, df, rest)
+    else:
+        return "Usage: svy: summarize [cols] | svy: ols y ~ x1 + x2 | svy: logit y ~ x1 + x2"
+
+
+def _svy_summarize(session: Session, df, args: str) -> str:
+    """Weighted summary statistics."""
+    cols = args.split() if args.strip() else [c for c in df.columns if df[c].dtype in NUMERIC_DTYPES]
+    if not cols:
+        return "No numeric columns to summarize."
+
+    from openstat.stats.survey import weighted_summary
+    return weighted_summary(df, cols, session._svy_weight_var)
+
+
+def _svy_ols(session: Session, df, args: str) -> str:
+    """Weighted OLS regression."""
+    try:
+        dep, indeps = parse_formula(args)
+    except ParseError as e:
+        return f"Formula error: {e}"
+
+    try:
+        from openstat.stats.survey import fit_weighted_ols
+        result, raw = fit_weighted_ols(
+            df, dep, indeps, session._svy_weight_var,
+            session._svy_strata_var, session._svy_psu_var,
+        )
+
+        session._last_model = raw
+        session._last_model_vars = (dep, indeps)
+        session._last_fit_result = result
+        session._last_fit_kwargs = {"survey": True}
+
+        md = result.to_markdown() if hasattr(result, "to_markdown") else ""
+        session.results.append(ModelResult(
+            name="Svy: OLS", formula=result.formula,
+            table=md, details={
+                "n_obs": result.n_obs,
+                "params": dict(result.params),
+                "r_squared": result.r_squared,
+            },
+        ))
+
+        output = result.summary_table()
+        if result.warnings:
+            output += "\n" + "\n".join(result.warnings)
+        return output
+    except Exception as e:
+        return friendly_error(e, "svy: ols")
+
+
+def _svy_logit(session: Session, df, args: str) -> str:
+    """Weighted logistic regression."""
+    try:
+        dep, indeps = parse_formula(args)
+    except ParseError as e:
+        return f"Formula error: {e}"
+
+    try:
+        from openstat.stats.survey import fit_weighted_logit
+        result, raw = fit_weighted_logit(df, dep, indeps, session._svy_weight_var)
+
+        session._last_model = raw
+        session._last_model_vars = (dep, indeps)
+        session._last_fit_result = result
+        session._last_fit_kwargs = {"survey": True}
+
+        md = result.to_markdown() if hasattr(result, "to_markdown") else ""
+        session.results.append(ModelResult(
+            name="Svy: Logit", formula=result.formula,
+            table=md, details={
+                "n_obs": result.n_obs,
+                "params": dict(result.params),
+            },
+        ))
+
+        output = result.summary_table()
+        if result.warnings:
+            output += "\n" + "\n".join(result.warnings)
+        return output
+    except Exception as e:
+        return friendly_error(e, "svy: logit")