framework-py 0.1.0.dev0__py3-none-any.whl

framework/outputs.py

@@ -0,0 +1,324 @@
+"""Saving results: tables, figures, models, reports, and rendered notebooks.
+
+Each save writes the file to the project's output folder for its kind and
+records it in the ``results`` table of framework.db (named by file name, with
+the file's hash), the same way R does. Sensitive projects keep public and
+private outputs in separate folders.
+"""
+
+from __future__ import annotations
+
+import pickle
+import shutil
+import subprocess
+import tempfile
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+
+from framework import db
+from framework._log import inform
+from framework._paths import output_directory
+from framework.errors import FrameworkError, UnsupportedFormatError
+from framework.integrity import file_hash
+from framework.settings import settings_read
+
+_TABLE_FORMATS = {"csv": ".csv", "parquet": ".parquet", "xlsx": ".xlsx"}
+_FIGURE_FORMATS = ("png", "pdf", "svg", "jpg", "jpeg", "tiff")
+_NOTEBOOK_FORMATS = {"html": ".html", "pdf": ".pdf", "docx": ".docx"}
+_RESULT_TYPES = ("table", "figure", "model", "report", "notebook")
+
+
+def _folder(kind: str, public: bool) -> Path:
+    try:
+        config = settings_read()
+    except FrameworkError:
+        config = {}
+    folder = output_directory(config, kind, public)
+    if not folder.is_dir():
+        folder.mkdir(parents=True, exist_ok=True)
+        inform(f"Creating {kind} directory: {folder}")
+    return folder
+
+
+def _check_overwrite(path: Path, overwrite: bool) -> None:
+    if path.exists() and not overwrite:
+        raise FrameworkError(f"File already exists: {path}. Use overwrite=True to replace.")
+
+
+def _record(path: Path, type: str, public: bool) -> None:
+    """Record a saved output in the results table, as R's ``.save_result()`` does."""
+    name = path.name
+    digest = file_hash(path) if path.is_file() else None
+    now = db.timestamp()
+    with db.connect() as connection:
+        if connection.execute("SELECT id FROM results WHERE name = ?", (name,)).fetchone():
+            connection.execute(
+                "UPDATE results SET type = ?, public = ?, comment = ?, hash = ?, updated_at = ? "
+                "WHERE name = ?",
+                (type, int(public), None, digest, now, name),
+            )
+        else:
+            connection.execute(
+                "INSERT INTO results (name, type, public, blind, comment, hash, created_at, "
+                "updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+                (name, type, int(public), 0, None, digest, now, now),
+            )
+
+
+def save_table(
+    data: pd.DataFrame,
+    name: str,
+    format: str = "csv",
+    public: bool = False,
+    overwrite: bool = True,
+    **kwargs: Any,
+) -> Path:
+    """Save a table to the tables output folder.
+
+    Args:
+        data: The table.
+        name: File name without extension.
+        format: ``"csv"``, ``"parquet"``, or ``"xlsx"`` (needs the ``excel`` extra).
+        public: In a sensitive project, save to the public tables folder.
+        overwrite: Replace an existing file.
+        **kwargs: Passed to the pandas writer.
+
+    Returns:
+        The saved file.
+    """
+    if format == "rds":
+        raise UnsupportedFormatError(".rds tables are R-only. Use csv, parquet, or xlsx.")
+    if format not in _TABLE_FORMATS:
+        raise FrameworkError(f"format must be one of {', '.join(_TABLE_FORMATS)}")
+    path = _folder("tables", public) / f"{name}{_TABLE_FORMATS[format]}"
+    _check_overwrite(path, overwrite)
+    if format == "csv":
+        data.to_csv(path, index=False, **{"lineterminator": "\n", **kwargs})
+    elif format == "parquet":
+        data.to_parquet(path, index=False, **kwargs)
+    else:
+        try:
+            import openpyxl  # noqa: F401
+        except ImportError as error:
+            raise FrameworkError(
+                "xlsx tables need openpyxl. Install with: pip install framework-py[excel]"
+            ) from error
+        data.to_excel(path, index=False, **kwargs)
+    inform(f"Saved table to {path}")
+    _record(path, "table", public)
+    return path
+
+
+def _current_figure() -> Any:
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError as error:
+        raise FrameworkError(
+            "save_figure(None, ...) saves the current matplotlib figure, "
+            "but matplotlib is not installed."
+        ) from error
+    return plt.gcf()
+
+
+def save_figure(
+    plot: Any,
+    name: str,
+    format: str = "png",
+    width: float = 8,
+    height: float = 6,
+    dpi: int = 300,
+    public: bool = False,
+    overwrite: bool = True,
+    **kwargs: Any,
+) -> Path:
+    """Save a figure to the figures output folder.
+
+    Args:
+        plot: A matplotlib or seaborn figure (anything with ``savefig()``), a
+            plotnine plot (anything with ``save()``), or None for the current
+            matplotlib figure.
+        name: File name without extension.
+        format: png, pdf, svg, jpg, jpeg, or tiff.
+        width: Width in inches.
+        height: Height in inches.
+        dpi: Resolution for raster formats.
+        public: In a sensitive project, save to the public figures folder.
+        overwrite: Replace an existing file.
+        **kwargs: Passed to ``savefig()`` or ``save()``.
+
+    Returns:
+        The saved file.
+    """
+    if format not in _FIGURE_FORMATS:
+        raise FrameworkError(f"format must be one of {', '.join(_FIGURE_FORMATS)}")
+    if width < 0.1 or height < 0.1 or dpi < 1:
+        raise FrameworkError("width and height must be at least 0.1 and dpi at least 1")
+    figure = _current_figure() if plot is None else plot
+    path = _folder("figures", public) / f"{name}.{format}"
+    _check_overwrite(path, overwrite)
+
+    if hasattr(figure, "savefig"):
+        target = getattr(figure, "figure", figure)
+        if hasattr(target, "set_size_inches"):
+            target.set_size_inches(width, height)
+        figure.savefig(path, dpi=dpi, format=format, **kwargs)
+    elif hasattr(figure, "save"):
+        figure.save(path, width=width, height=height, dpi=dpi, units="in", verbose=False, **kwargs)
+    else:
+        raise FrameworkError(
+            "plot must have a savefig() method (matplotlib, seaborn) or a save() method "
+            "(plotnine), or be None to save the current matplotlib figure"
+        )
+    inform(f"Saved figure to {path}")
+    _record(path, "figure", public)
+    return path
+
+
+def save_model(model: Any, name: str, public: bool = False, overwrite: bool = True) -> Path:
+    """Save a fitted model to the models output folder with pickle (``<name>.pkl``)."""
+    path = _folder("models", public) / f"{name}.pkl"
+    _check_overwrite(path, overwrite)
+    with path.open("wb") as handle:
+        pickle.dump(model, handle)
+    inform(f"Saved model to {path}")
+    _record(path, "model", public)
+    return path
+
+
+def save_report(
+    file: str | Path,
+    name: str | None = None,
+    public: bool = False,
+    overwrite: bool = True,
+    move: bool = False,
+) -> Path:
+    """Copy (or move) a finished report into the reports output folder.
+
+    Args:
+        file: The report file.
+        name: New file name without extension. Default: keep the name.
+        public: In a sensitive project, save to the public reports folder.
+        overwrite: Replace an existing file.
+        move: Move the file instead of copying it.
+    """
+    source = Path(file)
+    if not source.is_file():
+        raise FrameworkError(f"File not found: {source}")
+    folder = _folder("reports", public)
+    destination = folder / (source.name if name is None else f"{name}{source.suffix}")
+    _check_overwrite(destination, overwrite)
+    if move:
+        shutil.move(source, destination)
+        inform(f"Moved report to {destination}")
+    else:
+        shutil.copyfile(source, destination)
+        inform(f"Saved report to {destination}")
+    _record(destination, "report", public)
+    return destination
+
+
+def save_notebook(
+    file: str | Path,
+    name: str | None = None,
+    format: str = "html",
+    public: bool = False,
+    overwrite: bool = True,
+    embed_resources: bool = True,
+    quarto_args: Sequence[str] = (),
+) -> Path:
+    """Render a Quarto notebook or Jupyter notebook and save the result.
+
+    Rendering uses the ``quarto`` command, which must be installed.
+
+    Args:
+        file: A ``.qmd`` or ``.ipynb`` file.
+        name: Output file name without extension. Default: the notebook's name.
+        format: ``"html"``, ``"pdf"``, or ``"docx"``.
+        public: In a sensitive project, save to the public notebooks folder.
+        overwrite: Replace an existing file.
+        embed_resources: For HTML, embed images and styles in one file.
+        quarto_args: Extra arguments for ``quarto render``.
+    """
+    source = Path(file)
+    if not source.is_file() or source.suffix.lower() not in (".qmd", ".ipynb"):
+        raise FrameworkError(f"save_notebook() needs an existing .qmd or .ipynb file: {source}")
+    if format not in _NOTEBOOK_FORMATS:
+        raise FrameworkError(f"format must be one of {', '.join(_NOTEBOOK_FORMATS)}")
+    quarto = shutil.which("quarto")
+    if quarto is None:
+        raise FrameworkError(
+            "Quarto not found. Install it from https://quarto.org/docs/get-started/"
+        )
+
+    folder = _folder("notebooks", public)
+    destination = folder / f"{name or source.stem}{_NOTEBOOK_FORMATS[format]}"
+    _check_overwrite(destination, overwrite)
+
+    with tempfile.TemporaryDirectory(prefix="save_notebook_") as temp:
+        command = [quarto, "render", str(source.resolve()), "--output-dir", temp, "--to", format]
+        if embed_resources and format == "html":
+            command.append("--embed-resources")
+        command.extend(quarto_args)
+        inform(f"Rendering {source.name}...")
+        result = subprocess.run(command, capture_output=True, text=True, check=False)
+        if result.returncode != 0:
+            raise FrameworkError(f"Quarto render failed\n{result.stdout}{result.stderr}")
+        outputs = sorted(Path(temp).glob(f"*{_NOTEBOOK_FORMATS[format]}"))
+        if not outputs:
+            raise FrameworkError(f"No {format} output found after rendering")
+        shutil.copyfile(outputs[0], destination)
+    inform(f"Saved notebook to {destination}")
+    _record(destination, "notebook", public)
+    return destination
+
+
+def result_list(type: str | None = None, public: bool | None = None) -> pd.DataFrame:
+    """Saved results, newest first.
+
+    Columns: name, type, public, comment, hash, created_at, and updated_at.
+    Timestamps are local ``%Y-%m-%d %H:%M:%S`` strings, as in R.
+
+    Args:
+        type: Only this kind: table, figure, model, report, or notebook.
+        public: Only public (True) or private (False) results.
+    """
+    if type is not None and type not in _RESULT_TYPES:
+        raise FrameworkError(f"type must be one of {', '.join(_RESULT_TYPES)}")
+    columns = ["name", "type", "public", "comment", "hash", "created_at", "updated_at"]
+    query = (
+        "SELECT name, type, public, comment, hash, created_at, updated_at "
+        "FROM results WHERE deleted_at IS NULL"
+    )
+    params: list[Any] = []
+    if type is not None:
+        query += " AND type = ?"
+        params.append(type)
+    if public is not None:
+        query += " AND public = ?"
+        params.append(int(public))
+    with db.connect() as connection:
+        rows = connection.execute(query, params).fetchall()
+
+    records = []
+    for row in rows:
+        created = db.parse_timestamp(row["created_at"])
+        updated = db.parse_timestamp(row["updated_at"])
+        records.append(
+            {
+                "name": row["name"],
+                "type": row["type"],
+                "public": bool(row["public"]),
+                "comment": row["comment"],
+                "hash": row["hash"],
+                "created_at": created.strftime("%Y-%m-%d %H:%M:%S") if created else None,
+                "updated_at": updated.strftime("%Y-%m-%d %H:%M:%S") if updated else None,
+                "_sort": updated.timestamp() if updated else float("-inf"),
+            }
+        )
+    records.sort(key=lambda record: record["_sort"], reverse=True)
+    for record in records:
+        del record["_sort"]
+    return pd.DataFrame(records, columns=columns)

framework/packages.py

@@ -0,0 +1,276 @@
+"""Project packages, kept in ``pyproject.toml`` and locked with uv.
+
+R's ``packages_*()`` functions manage renv. In Python the dependencies live in
+``pyproject.toml``, ``uv.lock`` pins them, and these functions call uv. When
+uv is not installed, ``package_add()`` and ``packages_install()`` fall back to
+pip in the current Python environment. Locking needs uv.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import shutil
+import subprocess
+import sys
+import tomllib
+from dataclasses import dataclass, field
+from importlib.metadata import PackageNotFoundError, version
+from pathlib import Path
+
+import pandas as pd
+
+from framework._log import inform, show_messages
+from framework.errors import FrameworkError
+from framework.project import project_root
+
+_NAME = re.compile(r"^\s*([A-Za-z0-9][A-Za-z0-9._-]*)")
+
+
+def _run(args: list[str], root: Path) -> int:
+    """Run a command in the project folder, showing its output. Returns its exit status."""
+    return subprocess.run(args, cwd=root, check=False).returncode
+
+
+def _uv() -> str | None:
+    return shutil.which("uv")
+
+
+def _require_uv(task: str) -> str:
+    uv = _uv()
+    if uv is None:
+        raise FrameworkError(
+            f"{task} needs uv, which is not installed. "
+            "Install it from https://docs.astral.sh/uv/ and try again."
+        )
+    return uv
+
+
+def _pyproject(root: Path) -> Path:
+    path = root / "pyproject.toml"
+    if not path.is_file():
+        raise FrameworkError(f"No pyproject.toml in {root}. Python projects list packages there.")
+    return path
+
+
+def _dependencies(root: Path) -> list[str]:
+    data = tomllib.loads(_pyproject(root).read_text(encoding="utf-8"))
+    dependencies = (data.get("project") or {}).get("dependencies") or []
+    return [str(item) for item in dependencies]
+
+
+def _package_name(requirement: str) -> str:
+    match = _NAME.match(requirement)
+    if match is None:
+        raise ValueError(f"Not a package requirement: {requirement!r}")
+    return match.group(1)
+
+
+def _normalize(name: str) -> str:
+    return re.sub(r"[-_.]+", "-", name).lower()
+
+
+def _installed_version(name: str) -> str | None:
+    try:
+        return version(name)
+    except PackageNotFoundError:
+        return None
+
+
+def _check(status: int, what: str) -> None:
+    if status != 0:
+        raise FrameworkError(f"{what} failed (exit status {status})")
+
+
+def _add_to_pyproject(path: Path, requirement: str) -> bool:
+    """Add a requirement to ``[project] dependencies`` without uv. Returns whether it changed.
+
+    The file is edited as text so comments and layout stay as they are.
+    """
+    text = path.read_text(encoding="utf-8")
+    name = _normalize(_package_name(requirement))
+    if any(_normalize(_package_name(d)) == name for d in _dependencies(path.parent)):
+        return False
+    project = re.search(r"(?ms)^\[project\]\s*$(.*?)(?=^\[|\Z)", text)
+    block = re.search(r"(?ms)^dependencies\s*=\s*\[(.*?)\]", project.group(1)) if project else None
+    if project is None or block is None:
+        raise FrameworkError(
+            "Could not find the dependencies list under [project] in pyproject.toml. "
+            f"Add {requirement!r} there yourself, or install uv."
+        )
+    start = project.start(1) + block.start(1)
+    end = project.start(1) + block.end(1)
+    items = block.group(1)
+    entry = json.dumps(requirement)
+    if items.strip():
+        body = items.rstrip()
+        separator = "" if body.endswith(",") else ","
+        updated = f"{body}{separator}\n    {entry},\n"
+    else:
+        updated = f"\n    {entry},\n"
+    path.write_text(text[:start] + updated + text[end:], encoding="utf-8")
+    return True
+
+
+# Public API -----------------------------------------------------------------
+
+
+def package_add(requirement: str, *, install: bool = True, dev: bool = False) -> None:
+    """Add a package to the project's dependencies, and install it.
+
+    Args:
+        requirement: A package name, optionally with a version, such as
+            ``"pandas"`` or ``"statsmodels>=0.14"``.
+        install: Install it now. With uv, False only updates
+            ``pyproject.toml`` and ``uv.lock``.
+        dev: Add it to the ``dev`` dependency group (uv only).
+    """
+    show_messages()
+    root = project_root()
+    path = _pyproject(root)
+    _package_name(requirement)
+    uv = _uv()
+    if uv is not None:
+        args = [uv, "add", requirement]
+        if dev:
+            args.append("--dev")
+        if not install:
+            args.append("--no-sync")
+        _check(_run(args, root), f"uv add {requirement}")
+        inform(f"✓ Added {requirement}")
+        return
+    if dev:
+        raise FrameworkError("Dev dependencies need uv. Install it from https://docs.astral.sh/uv/")
+    if _add_to_pyproject(path, requirement):
+        inform(f"✓ Added {requirement} to pyproject.toml")
+    else:
+        inform(f"{_package_name(requirement)} is already in pyproject.toml")
+    if install:
+        _check(_run([sys.executable, "-m", "pip", "install", requirement], root), "pip install")
+
+
+def packages_list() -> pd.DataFrame:
+    """The project's dependencies from ``pyproject.toml``, with installed versions.
+
+    Returns:
+        A DataFrame with columns package, requirement, and installed (None
+        when the package is not installed in the running environment).
+    """
+    root = project_root()
+    rows = [
+        {
+            "package": _package_name(requirement),
+            "requirement": requirement,
+            "installed": _installed_version(_package_name(requirement)),
+        }
+        for requirement in _dependencies(root)
+    ]
+    return pd.DataFrame(rows, columns=["package", "requirement", "installed"])
+
+
+def packages_install() -> None:
+    """Install the project's dependencies: ``uv sync``, or pip without uv."""
+    show_messages()
+    root = project_root()
+    uv = _uv()
+    if uv is not None:
+        _check(_run([uv, "sync"], root), "uv sync")
+        return
+    dependencies = _dependencies(root)
+    if dependencies:
+        _check(_run([sys.executable, "-m", "pip", "install", *dependencies], root), "pip install")
+
+
+def packages_restore() -> None:
+    """Install exactly the versions in ``uv.lock`` (``uv sync --frozen``)."""
+    show_messages()
+    root = project_root()
+    uv = _require_uv("Restoring packages")
+    if not (root / "uv.lock").is_file():
+        raise FrameworkError("No uv.lock to restore from. Run packages_snapshot() first.")
+    _check(_run([uv, "sync", "--frozen"], root), "uv sync --frozen")
+
+
+def packages_snapshot() -> None:
+    """Record the resolved versions in ``uv.lock`` (``uv lock``)."""
+    show_messages()
+    root = project_root()
+    _pyproject(root)
+    _check(_run([_require_uv("Locking packages"), "lock"], root), "uv lock")
+
+
+def packages_update(packages: str | list[str] | None = None) -> None:
+    """Upgrade locked versions, all of them or the named ones, and install them."""
+    show_messages()
+    root = project_root()
+    uv = _require_uv("Updating packages")
+    args = [uv, "lock"]
+    if packages is None:
+        args.append("--upgrade")
+    else:
+        for name in [packages] if isinstance(packages, str) else packages:
+            args += ["--upgrade-package", name]
+    _check(_run(args, root), "uv lock")
+    _check(_run([uv, "sync"], root), "uv sync")
+
+
+@dataclass
+class PackagesStatus:
+    """Whether the lock file and the environment match ``pyproject.toml``.
+
+    ``lock_current`` and ``synced`` are None when uv is not installed.
+    """
+
+    lock_current: bool | None
+    synced: bool | None
+    missing: list[str] = field(default_factory=list)
+
+    @property
+    def ok(self) -> bool:
+        return self.lock_current is not False and self.synced is not False and not self.missing
+
+
+def packages_status() -> PackagesStatus:
+    """Check that ``uv.lock`` is current and the environment matches it.
+
+    ``missing`` lists dependencies that are not installed in the running
+    Python environment.
+    """
+    root = project_root()
+    missing = [
+        _package_name(requirement)
+        for requirement in _dependencies(root)
+        if _installed_version(_package_name(requirement)) is None
+    ]
+    uv = _uv()
+    lock_current = synced = None
+    if uv is not None:
+
+        def quiet(*args: str) -> bool:
+            completed = subprocess.run(
+                [uv, *args],
+                cwd=root,
+                check=False,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+            return completed.returncode == 0
+
+        if (root / "uv.lock").is_file():
+            lock_current = quiet("lock", "--check")
+            synced = quiet("sync", "--check", "--frozen")
+        else:
+            lock_current = False
+    status = PackagesStatus(lock_current=lock_current, synced=synced, missing=missing)
+    lines = [
+        f"Lock file: {_describe(lock_current, 'current', 'out of date (run packages_snapshot())')}",
+        f"Environment: {_describe(synced, 'in sync', 'out of sync (run packages_restore())')}",
+    ]
+    if missing:
+        lines.append(f"Not installed here: {', '.join(missing)}")
+    print("\n".join(lines))
+    return status
+
+
+def _describe(value: bool | None, good: str, bad: str) -> str:
+    return "unknown (uv is not installed)" if value is None else good if value else bad

framework/project.py

@@ -0,0 +1,52 @@
+"""Finding the project root and its settings file."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from framework.errors import ProjectNotFoundError
+
+SETTINGS_FILES = ("settings.yml", "config.yml")
+
+# R stops after ten parent folders, so Python does too.
+_MAX_DEPTH = 10
+
+
+def _settings_in(folder: Path) -> Path | None:
+    for name in SETTINGS_FILES:
+        candidate = folder / name
+        if candidate.is_file():
+            return candidate
+    return None
+
+
+def find_settings_file(start: str | Path | None = None) -> Path | None:
+    """The project's settings file, searching ``start`` and up to ten parent folders.
+
+    ``settings.yml`` is preferred over the older ``config.yml``.
+    """
+    current = Path(start or Path.cwd()).resolve()
+    for _ in range(_MAX_DEPTH + 1):
+        found = _settings_in(current)
+        if found is not None:
+            return found
+        if current.parent == current:
+            return None
+        current = current.parent
+    return None
+
+
+def project_root(start: str | Path | None = None) -> Path:
+    """The root folder of the Framework project containing ``start`` (default: the current folder).
+
+    Raises:
+        ProjectNotFoundError: No settings.yml or config.yml was found.
+    """
+    found = find_settings_file(start)
+    if found is None:
+        where = Path(start or Path.cwd()).resolve()
+        raise ProjectNotFoundError(
+            f"No Framework project found at or above {where}. "
+            "A project has settings.yml (or config.yml) in its root folder."
+        )
+    return found.parent