overture-schema-codegen 0.1.1.dev0__py3-none-any.whl

overture/schema/codegen/__init__.py

@@ -0,0 +1 @@
+"""Code generator for Overture Schema Pydantic models."""

overture/schema/codegen/cli.py

@@ -0,0 +1,228 @@
+"""CLI entrypoint for schema code generation."""
+
+import json
+import logging
+from collections.abc import Mapping
+from pathlib import Path, PurePosixPath
+
+import click
+
+from overture.schema.cli.tag_options import build_selector, tag_selection_options
+from overture.schema.system.discovery import (
+    discover_models,
+    filter_models,
+    split_entry_point,
+)
+
+from .extraction.specs import ModelSpec, SupplementarySpec, TypeIdentity
+from .layout.module_layout import (
+    OUTPUT_ROOT,
+    compute_schema_root,
+    entry_point_module,
+)
+from .markdown.pipeline import generate_markdown_pages
+from .pyspark.pipeline import generate_pyspark_modules
+from .spec_discovery import extract_alias_spec, extract_model_spec
+
+log = logging.getLogger(__name__)
+
+__all__ = ["cli"]
+
+_OUTPUT_FORMATS = ("markdown", "pyspark")
+
+_FEATURE_FRONTMATTER = "---\nsidebar_position: 1\n---\n\n"
+
+
+def _write_output(
+    content: str,
+    output_dir: Path | None,
+    output_path: PurePosixPath,
+) -> None:
+    """Write content to a file under output_dir, or stdout."""
+    if output_dir:
+        file_path = output_dir / output_path
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        file_path.write_text(content)
+    else:
+        click.echo(content)
+        click.echo()  # separate entries with a blank line in stdout mode
+
+
+@click.group()
+def cli() -> None:
+    """Overture Schema code generator.
+
+    Generate documentation and code from Pydantic schema models.
+    """
+
+
+@cli.command("list")
+def list_models() -> None:
+    """List all discovered models."""
+    models = discover_models()
+    # Name every entry from its entry point, not the loaded object: a
+    # discriminated union loads as an `Annotated[...]` alias with no
+    # `__name__`, so `str(model)` would print the whole type expression.
+    names = []
+    for key in models:
+        _, class_name = split_entry_point(key.entry_point)
+        names.append(class_name)
+    for name in sorted(names):
+        click.echo(name)
+
+
+@cli.command()
+@click.option(
+    "--format",
+    "output_format",
+    required=True,
+    type=click.Choice(_OUTPUT_FORMATS),
+    help="Output format",
+)
+@tag_selection_options
+@click.option(
+    "--output-dir",
+    type=click.Path(path_type=Path),
+    default=None,
+    help="Write output files directly into this directory (default: stdout). "
+    "For pyspark, writes expression modules (*.py). "
+    "For markdown, writes theme subdirectories.",
+)
+@click.option(
+    "--test-output-dir",
+    type=click.Path(path_type=Path),
+    default=None,
+    help="Write test modules (test_*.py) into this directory (pyspark only).",
+)
+def generate(
+    output_format: str,
+    tags: tuple[str, ...],
+    filters: tuple[str, ...],
+    excludes: tuple[str, ...],
+    output_dir: Path | None,
+    test_output_dir: Path | None,
+) -> None:
+    """Generate code/docs from discovered models."""
+    if output_format != "pyspark" and test_output_dir is not None:
+        raise click.UsageError("--test-output-dir is only valid with --format pyspark")
+
+    all_models = discover_models()
+
+    models = filter_models(all_models, build_selector(tags, filters, excludes))
+
+    if output_dir:
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+    model_specs: list[ModelSpec] = [
+        spec
+        for key, entry in models.items()
+        if (spec := extract_model_spec(key, entry)) is not None
+    ]
+
+    if output_format == "pyspark":
+        _generate_pyspark(model_specs, output_dir, test_output_dir)
+    else:
+        # RootModel entry points yield no ModelSpec, so they document as
+        # named aliases -- reachable no other way, since a RootModel field
+        # unwraps to its bare shape and names no type.
+        external_specs: dict[TypeIdentity, SupplementarySpec] = {
+            alias.identity: alias
+            for entry in models.values()
+            if (alias := extract_alias_spec(entry)) is not None
+        }
+        module_paths = [entry_point_module(k.entry_point) for k in all_models]
+        schema_root = compute_schema_root(module_paths)
+        _generate_markdown(model_specs, schema_root, output_dir, external_specs)
+
+
+def _generate_markdown(
+    model_specs: list[ModelSpec],
+    schema_root: str,
+    output_dir: Path | None,
+    external_specs: Mapping[TypeIdentity, SupplementarySpec],
+) -> None:
+    """Generate markdown with directory layout and placement-aware links."""
+    pages = generate_markdown_pages(
+        model_specs, schema_root, external_specs=external_specs
+    )
+
+    for page in pages:
+        content = (
+            f"{_FEATURE_FRONTMATTER}{page.content}" if page.is_model else page.content
+        )
+        _write_output(content, output_dir, page.path)
+
+    if output_dir:
+        feature_paths = {page.path for page in pages if page.is_model}
+        all_paths = {page.path for page in pages}
+        _write_category_files(output_dir, all_paths, feature_paths)
+
+
+def _generate_pyspark(
+    model_specs: list[ModelSpec],
+    output_dir: Path | None,
+    test_output_dir: Path | None = None,
+) -> None:
+    """Generate PySpark validation modules.
+
+    Output is syntactically valid Python; we assume a code formatter runs
+    over the written directories afterwards to match existing conventions.
+    """
+    modules = generate_pyspark_modules(model_specs)
+    for mod in modules.source:
+        _write_output(mod.content, output_dir, mod.path)
+    if test_output_dir is not None:
+        for mod in modules.test:
+            _write_output(mod.content, test_output_dir, mod.path)
+
+
+def _ancestor_dirs(paths: set[PurePosixPath]) -> set[PurePosixPath]:
+    """Collect all ancestor directories for a set of file paths."""
+    dirs: set[PurePosixPath] = set()
+    for path in paths:
+        parent = path.parent
+        while parent != OUTPUT_ROOT:
+            dirs.add(parent)
+            parent = parent.parent
+    return dirs
+
+
+def _top_level_positions(
+    dirs: set[PurePosixPath],
+    feature_paths: set[PurePosixPath],
+) -> dict[PurePosixPath, int]:
+    """Assign sidebar positions: feature dirs first, then non-feature, both alphabetical."""
+    feature_dir_names = {p.parts[0] for p in feature_paths}
+    top_level = sorted(d for d in dirs if d.parent == OUTPUT_ROOT)
+    feature_dirs = [d for d in top_level if d.name in feature_dir_names]
+    non_feature_dirs = [d for d in top_level if d.name not in feature_dir_names]
+    return {d: i for i, d in enumerate(feature_dirs + non_feature_dirs, start=1)}
+
+
+def _write_category_files(
+    output_dir: Path,
+    all_paths: set[PurePosixPath],
+    feature_paths: set[PurePosixPath],
+) -> None:
+    """Write _category_.json files for Docusaurus sidebar navigation."""
+    dirs = _ancestor_dirs(all_paths)
+    positions = _top_level_positions(dirs, feature_paths)
+
+    for dir_path in sorted(dirs):
+        label = dir_path.name.replace("_", " ").title()
+        category: dict[str, object] = {"label": label}
+        if dir_path in positions:
+            category["position"] = positions[dir_path]
+
+        file_path = output_dir / dir_path / "_category_.json"
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        file_path.write_text(json.dumps(category, indent=2) + "\n")
+
+
+def main() -> None:
+    """Run the CLI entry point."""
+    cli()
+
+
+if __name__ == "__main__":
+    main()

overture/schema/codegen/extraction/docstring.py

@@ -0,0 +1,46 @@
+"""Docstring extraction and cleaning utilities."""
+
+import inspect
+from enum import Enum
+from typing import NewType
+
+__all__ = ["clean_docstring", "first_docstring_line", "is_custom_docstring"]
+
+
+# Probe auto-generated docstrings so we can distinguish them from explicit ones.
+# Both Enum and NewType generate default docstrings that vary by Python version;
+# capturing at import time adapts automatically if the format changes.
+class _DocstringProbeEnum(Enum):
+    pass
+
+
+_ENUM_DEFAULT_DOCSTRING = _DocstringProbeEnum.__doc__
+del _DocstringProbeEnum
+_NewtypeProbe = NewType("_NewtypeProbe", int)
+_NEWTYPE_DEFAULT_DOCSTRING = _NewtypeProbe.__doc__
+del _NewtypeProbe
+
+
+def clean_docstring(doc: str | None) -> str | None:
+    """Return cleaned docstring, or None if absent or whitespace-only."""
+    if not doc:
+        return None
+    cleaned = inspect.cleandoc(doc)
+    return cleaned or None
+
+
+def first_docstring_line(doc: str | None) -> str | None:
+    """Return the first line of a docstring, or None if absent."""
+    cleaned = clean_docstring(doc)
+    if not cleaned:
+        return None
+    return cleaned.split("\n")[0]
+
+
+def is_custom_docstring(doc: str | None, inherited_doc: str | None = None) -> bool:
+    """Check if a docstring was explicitly written, not auto-generated or inherited."""
+    return bool(doc) and doc not in (
+        _ENUM_DEFAULT_DOCSTRING,
+        _NEWTYPE_DEFAULT_DOCSTRING,
+        inherited_doc,
+    )

overture/schema/codegen/extraction/enum_extraction.py

@@ -0,0 +1,40 @@
+"""Enum extraction."""
+
+from enum import Enum
+
+from .docstring import clean_docstring, is_custom_docstring
+from .specs import EnumMemberSpec, EnumSpec
+
+__all__ = ["extract_enum"]
+
+
+def extract_enum(enum_class: type[Enum]) -> EnumSpec:
+    """Extract enum specification from an Enum class.
+
+    Handles both simple str Enums and DocumentedEnums where members
+    have per-value descriptions via the __doc__ attribute.
+    """
+    class_doc = enum_class.__doc__
+    description = clean_docstring(class_doc) if is_custom_docstring(class_doc) else None
+
+    members: list[EnumMemberSpec] = []
+    for member in enum_class:
+        member_doc = getattr(member, "__doc__", None)
+        member_description = (
+            member_doc if is_custom_docstring(member_doc, class_doc) else None
+        )
+
+        members.append(
+            EnumMemberSpec(
+                name=member.name,
+                value=str(member.value),
+                description=member_description,
+            )
+        )
+
+    return EnumSpec(
+        name=enum_class.__name__,
+        description=description,
+        members=members,
+        source_type=enum_class,
+    )

overture/schema/codegen/extraction/examples.py

@@ -0,0 +1,367 @@
+"""Load, validate, and flatten example data for schema documentation."""
+
+import logging
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel, TypeAdapter, ValidationError
+from pydantic.fields import FieldInfo
+
+from .model_extraction import resolve_field_alias
+from .type_analyzer import single_literal_value
+
+log = logging.getLogger(__name__)
+
+__all__ = [
+    "ExampleRecord",
+    "augment_missing_fields",
+    "flatten_model_instance",
+    "load_examples",
+    "validate_example",
+]
+
+# tomllib is stdlib from 3.11+; tomli is the backport for 3.10.
+try:
+    import tomllib  # type: ignore[import-not-found]
+except ModuleNotFoundError:
+    import tomli as tomllib  # type: ignore[import-not-found]
+
+
+@dataclass
+class ExampleRecord:
+    """A flattened example with field-value pairs in documentation order."""
+
+    rows: list[tuple[str, Any]]
+
+
+def _inject_literal_fields(
+    model_fields_dict: dict[str, FieldInfo], data: dict[str, Any]
+) -> dict[str, Any]:
+    """Inject single-value Literal field defaults missing from *data*.
+
+    Inspects *model_fields_dict* for fields with single-value `Literal`
+    annotations. For each field missing from *data*, injects the literal
+    value using the field's `validation_alias` (if set), falling back
+    to `alias`, then to the field name.
+
+    Returns a new dict; the original is not mutated.
+    """
+    result = data.copy()
+
+    for field_name, field_info in model_fields_dict.items():
+        key = resolve_field_alias(field_name, field_info)
+        if key in result:
+            continue
+
+        literal_value = single_literal_value(field_info.annotation)
+        if literal_value is not None:
+            result[key] = literal_value
+
+    return result
+
+
+def _known_field_keys(model_fields_dict: dict[str, FieldInfo]) -> frozenset[str]:
+    """Alias-resolved field keys from a model_fields dict."""
+    return frozenset(
+        resolve_field_alias(name, info) for name, info in model_fields_dict.items()
+    )
+
+
+def _strip_null_unknown_fields(
+    data: dict[str, Any], known_keys: frozenset[str]
+) -> dict[str, Any]:
+    """Drop null-valued fields not in *known_keys*.
+
+    For discriminated unions, *known_keys* contains only common base
+    fields.  Variant-specific null fields from other arms (present in
+    flat parquet schemas) are stripped so the selected arm's validator
+    doesn't reject them as unknown extras.
+
+    Non-null fields are always kept so the arm's own validator can
+    accept or reject them normally.
+    """
+    return {k: v for k, v in data.items() if v is not None or k in known_keys}
+
+
+def validate_example(
+    validation_type: object,
+    raw: dict[str, Any],
+    *,
+    model_fields: dict[str, FieldInfo] | None = None,
+) -> BaseModel:
+    """Validate example data against a model or union type.
+
+    Returns the validated model instance. Preprocesses *raw* data by:
+    1. Injecting missing Literal fields for validation (if model_fields provided)
+    2. Stripping null-valued fields not in *model_fields* (handles
+       flat-schema examples from discriminated unions)
+    """
+    if model_fields is None:
+        if isinstance(validation_type, type) and issubclass(validation_type, BaseModel):
+            model_fields = validation_type.model_fields
+        else:
+            model_fields = {}
+
+    known_keys = _known_field_keys(model_fields)
+    preprocessed = _inject_literal_fields(model_fields, raw)
+    preprocessed = _strip_null_unknown_fields(preprocessed, known_keys)
+    result: object = TypeAdapter(validation_type).validate_python(preprocessed)
+    if not isinstance(result, BaseModel):
+        raise TypeError(f"Expected BaseModel instance, got {type(result).__name__}")
+    return result
+
+
+def extract_base_field(key: str) -> str:
+    """Extract the top-level field name from a flattened key.
+
+    >>> extract_base_field("sources[0].dataset")
+    'sources'
+    >>> extract_base_field("names.primary")
+    'names'
+    >>> extract_base_field("id")
+    'id'
+    """
+    if "[" in key:
+        return key.split("[")[0]
+    if "." in key:
+        return key.split(".")[0]
+    return key
+
+
+def order_example_rows(
+    flat_rows: list[tuple[str, Any]],
+    field_names: list[str],
+) -> list[tuple[str, Any]]:
+    """Order flattened rows by field position in documentation.
+
+    Sorts by position of base field name in *field_names*.
+    Fields with the same base maintain their original order (stable sort).
+    Unknown fields sort to end.
+    """
+    position = {name: i for i, name in enumerate(field_names)}
+    sentinel = len(field_names)
+
+    def sort_key(row: tuple[str, Any]) -> int:
+        return position.get(extract_base_field(row[0]), sentinel)
+
+    return sorted(flat_rows, key=sort_key)
+
+
+def _structured_fields(value: object) -> list[tuple[str, Any]] | None:
+    """Extract named fields from `__slots__`-based types like BBox.
+
+    Returns a list of `(name, value)` pairs for types that expose
+    public properties backed by private slots (`_name` -> `name`).
+    Returns `None` for types without this pattern.
+    """
+    cls = type(value)
+    slots = getattr(cls, "__slots__", ())
+    if not slots:
+        return None
+    fields: list[tuple[str, Any]] = []
+    for slot in slots:
+        attr = slot.lstrip("_")
+        if attr != slot and isinstance(getattr(cls, attr, None), property):
+            fields.append((attr, getattr(value, attr)))
+    return fields if len(fields) >= 2 else None
+
+
+def _needs_recursion(items: list[Any]) -> bool:
+    """Check whether list items contain models or nested lists."""
+    return bool(items) and isinstance(items[0], (BaseModel, list))
+
+
+def _flatten_list_items(key: str, items: list[Any]) -> list[tuple[str, Any]]:
+    """Flatten list items, recursing into BaseModel and nested list items.
+
+    Returns the list as a single leaf value when no items need recursion.
+    Pydantic model fields produce homogeneous lists, so the first item's
+    type determines the flattening strategy.
+    """
+    if not _needs_recursion(items):
+        return [(key, items)]
+    rows: list[tuple[str, Any]] = []
+    for i, item in enumerate(items):
+        if isinstance(item, BaseModel):
+            rows.extend(flatten_model_instance(item, f"{key}[{i}]."))
+        elif isinstance(item, list):
+            rows.extend(_flatten_list_items(f"{key}[{i}]", item))
+        else:
+            rows.append((f"{key}[{i}]", item))
+    return rows
+
+
+def flatten_model_instance(
+    instance: BaseModel,
+    prefix: str = "",
+) -> list[tuple[str, Any]]:
+    """Flatten a Pydantic model instance to dot-notation key-value pairs.
+
+    Walks model fields recursively. BaseModel values recurse with dot
+    notation, lists of BaseModel recurse with bracket notation, and
+    everything else (dicts, primitives, None) is a leaf value.
+
+    Parameters
+    ----------
+    instance
+        The Pydantic model instance to flatten.
+    prefix
+        Dot-notation prefix accumulated from parent fields.
+
+    Returns
+    -------
+    list[tuple[str, Any]]
+        Flattened key-value pairs in field declaration order.
+    """
+    rows: list[tuple[str, Any]] = []
+    for field_name, field_info in type(instance).model_fields.items():
+        key = resolve_field_alias(field_name, field_info)
+        value = getattr(instance, field_name)
+        full_key = f"{prefix}{key}" if prefix else key
+
+        if isinstance(value, BaseModel):
+            rows.extend(flatten_model_instance(value, f"{full_key}."))
+        elif isinstance(value, list):
+            rows.extend(_flatten_list_items(full_key, value))
+        elif (sub_fields := _structured_fields(value)) is not None:
+            for name, v in sub_fields:
+                rows.append((f"{full_key}.{name}", v))
+        else:
+            rows.append((full_key, value))
+    return rows
+
+
+def augment_missing_fields(
+    rows: list[tuple[str, Any]],
+    field_names: list[str],
+) -> list[tuple[str, Any]]:
+    """Add (name, None) entries for fields absent from *rows*.
+
+    Compares base field names (via `extract_base_field`) against
+    *field_names*. Fields in *field_names* not represented in *rows*
+    are appended as `(name, None)`. Handles dot-notation and bracket-
+    notation keys correctly.
+
+    Parameters
+    ----------
+    rows
+        Flattened key-value pairs from a concrete model instance.
+    field_names
+        Merged field name list from the union spec.
+
+    Returns
+    -------
+    list[tuple[str, Any]]
+        Original rows with (name, None) entries appended for absent fields.
+    """
+    present = {extract_base_field(key) for key, _ in rows}
+    augmented = list(rows)
+    for name in field_names:
+        if name not in present:
+            augmented.append((name, None))
+    return augmented
+
+
+def load_examples_from_toml(
+    pyproject_path: Path,
+    model_name: str,
+) -> list[dict[str, Any]]:
+    """Load `[examples.<model_name>]` from a pyproject.toml file."""
+    with pyproject_path.open("rb") as f:
+        data = tomllib.load(f)
+
+    examples: dict[str, list[dict[str, Any]]] = data.get("examples", {})
+    return examples.get(model_name, [])
+
+
+def resolve_pyproject_path(model_class: type) -> Path | None:
+    """Find pyproject.toml by walking up from the model's module location."""
+    module_name = getattr(model_class, "__module__", None)
+    if not module_name:
+        return None
+
+    module = sys.modules.get(module_name)
+    if not module:
+        return None
+
+    module_file = getattr(module, "__file__", None)
+    if not module_file:
+        return None
+
+    # Walk up from module directory
+    current = Path(module_file).parent
+    while current != current.parent:  # Stop at filesystem root
+        pyproject = current / "pyproject.toml"
+        if pyproject.exists():
+            return pyproject
+        current = current.parent
+
+    return None
+
+
+def load_examples(
+    validation_type: object,
+    model_name: str,
+    field_names: list[str],
+    *,
+    pyproject_source: type | None = None,
+    model_fields: dict[str, FieldInfo] | None = None,
+) -> list[ExampleRecord]:
+    """Load examples for a model, flattened and ordered by *field_names*.
+
+    Validates each example against the validation type. Invalid examples
+    are skipped with a warning logged. Returns an empty list on any failure
+    (missing file, missing section, parse error).
+
+    Parameters
+    ----------
+    validation_type : type[BaseModel] | object
+        Model class or union alias to validate against.
+    model_name : str
+        Name of the model to load examples for.
+    field_names : list[str]
+        List of field names for ordering output.
+    pyproject_source : type or None
+        Type to use for finding pyproject.toml. If None,
+        uses validation_type if it's a class.
+    model_fields : dict[str, FieldInfo] or None
+        Field info dict for Literal injection. If None, infers
+        from validation_type if it's a BaseModel class.
+    """
+    source_type = pyproject_source if pyproject_source is not None else validation_type
+    if not isinstance(source_type, type):
+        return []
+
+    pyproject_path = resolve_pyproject_path(source_type)
+    if not pyproject_path:
+        return []
+
+    try:
+        raw_examples = load_examples_from_toml(pyproject_path, model_name)
+    except (OSError, tomllib.TOMLDecodeError):
+        log.debug("Failed to load examples for %s", model_name, exc_info=True)
+        return []
+
+    if not raw_examples:
+        return []
+
+    records = []
+    for raw in raw_examples:
+        try:
+            instance = validate_example(validation_type, raw, model_fields=model_fields)
+        except ValidationError as e:
+            log.warning(
+                "Skipping invalid example for %s in %s: %s",
+                model_name,
+                pyproject_path,
+                e,
+            )
+            continue
+        flat_rows = flatten_model_instance(instance)
+        flat_rows = augment_missing_fields(flat_rows, field_names)
+        ordered_rows = order_example_rows(flat_rows, field_names)
+        records.append(ExampleRecord(rows=ordered_rows))
+
+    return records