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

overture/schema/codegen/extraction/type_registry.py

@@ -0,0 +1,137 @@
+"""Type registry mapping Python types to target representations."""
+
+from dataclasses import dataclass
+from typing import Literal
+
+from .field import FieldShape
+from .field_walk import newtype_name, terminal_primitive
+
+__all__ = [
+    "SparkCategory",
+    "TypeMapping",
+    "PRIMITIVE_TYPES",
+    "get_type_mapping",
+    "is_semantic_newtype",
+    "primitive_spark_category",
+    "resolve_type_name",
+]
+
+SparkCategory = Literal["string", "int", "float", "bool", "other"]
+
+
+@dataclass(frozen=True)
+class TypeMapping:
+    """Maps a type to its representation in different targets."""
+
+    markdown: str
+    spark: str | None = None
+
+
+PRIMITIVE_TYPES: dict[str, TypeMapping] = {
+    # Signed integers
+    "int8": TypeMapping(markdown="int8", spark="IntegerType()"),
+    "int16": TypeMapping(markdown="int16", spark="IntegerType()"),
+    "int32": TypeMapping(markdown="int32", spark="IntegerType()"),
+    "int64": TypeMapping(markdown="int64", spark="LongType()"),
+    # Unsigned integers
+    "uint8": TypeMapping(markdown="uint8", spark="IntegerType()"),
+    "uint16": TypeMapping(markdown="uint16", spark="IntegerType()"),
+    "uint32": TypeMapping(markdown="uint32", spark="IntegerType()"),
+    # Floating point
+    "float32": TypeMapping(markdown="float32", spark="FloatType()"),
+    "float64": TypeMapping(markdown="float64", spark="DoubleType()"),
+    # Basic types
+    "str": TypeMapping(markdown="string", spark="StringType()"),
+    "bool": TypeMapping(markdown="boolean", spark="BooleanType()"),
+    # Python builtins (aliases to their portable equivalents)
+    "int": TypeMapping(markdown="int64", spark="LongType()"),
+    "float": TypeMapping(markdown="float64", spark="DoubleType()"),
+    # Geometry types
+    "Geometry": TypeMapping(markdown="geometry", spark="BinaryType()"),
+    "BBox": TypeMapping(markdown="bbox"),
+}
+
+
+def is_semantic_newtype(shape: FieldShape) -> bool:
+    """Whether a shape's outermost NewType should be displayed by name.
+
+    Returns True for unregistered NewTypes (HexColor, Sources) and
+    NewTypes that wrap a different base type (FeatureVersion wrapping
+    int32, Id wrapping NoWhitespaceString). Returns False for
+    registered primitives (int32, Geometry).
+    """
+    nt_name = newtype_name(shape)
+    if nt_name is None:
+        return False
+    terminal = terminal_primitive(shape)
+    if terminal is None:
+        return True
+    if nt_name != terminal.base_type:
+        return True
+    return get_type_mapping(terminal.base_type) is None
+
+
+def get_type_mapping(type_name: str) -> TypeMapping | None:
+    """Look up a type mapping by name.
+
+    Accepts portable type names (`int32`, `str`, `Geometry`) and Python
+    builtin names (`int` -> int64, `float` -> float64).
+    """
+    return PRIMITIVE_TYPES.get(type_name)
+
+
+# BinaryType() is intentionally absent: geometry maps to BinaryType() in
+# PRIMITIVE_TYPES but falls through to "other" here, not a numeric/string/bool scalar.
+_SPARK_TYPE_CATEGORIES: dict[str, SparkCategory] = {
+    "StringType()": "string",
+    "IntegerType()": "int",
+    "LongType()": "int",
+    "FloatType()": "float",
+    "DoubleType()": "float",
+    "BooleanType()": "bool",
+}
+
+
+def primitive_spark_category(base_type: str) -> SparkCategory:
+    """Return the Spark category for a primitive base type name.
+
+    Parameters
+    ----------
+    base_type
+        A primitive type name (`"int32"`, `"float64"`, `"bool"`, `"str"`, ...).
+
+    Returns
+    -------
+    SparkCategory
+        `"string"` for string-valued types, `"int"` for integer types,
+        `"float"` for floating-point types, `"bool"` for boolean types,
+        `"other"` for binary, geometry, or unregistered types. Unknown
+        types fall back to `"other"`, preserving string-default behavior
+        for any future unregistered type.
+    """
+    mapping = get_type_mapping(base_type)
+    if mapping is None or mapping.spark is None:
+        return "other"
+    return _SPARK_TYPE_CATEGORIES.get(mapping.spark, "other")
+
+
+def resolve_type_name(shape: FieldShape) -> str:
+    """Resolve a shape to its markdown base type name string.
+
+    Looks up the terminal scalar's `base_type` in the registry first,
+    falling back to `source_type.__name__`. Semantic NewTypes wrapping
+    unregistered types resolve to the underlying class name (e.g.
+    `Sources` wrapping `SourceItem` -> `SourceItem`).
+    """
+    terminal = terminal_primitive(shape)
+    if terminal is None:
+        return "?"
+    mapping = get_type_mapping(terminal.base_type)
+    if mapping is None and terminal.source_type is not None:
+        mapping = get_type_mapping(terminal.source_type.__name__)
+    if mapping is not None:
+        return mapping.markdown
+
+    if newtype_name(shape) and terminal.source_type is not None:
+        return terminal.source_type.__name__
+    return terminal.base_type

overture/schema/codegen/extraction/union_extraction.py

@@ -0,0 +1,270 @@
+"""Union extraction and discriminator handling."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from enum import Enum
+from typing import Annotated, get_args, get_origin
+
+from pydantic import BaseModel
+from pydantic.fields import FieldInfo
+
+from overture.schema.system.feature import resolve_discriminator_field_name
+
+from .field import (
+    AnyScalar,
+    ArrayOf,
+    FieldShape,
+    LiteralScalar,
+    MapOf,
+    ModelRef,
+    NewTypeShape,
+    Primitive,
+    UnionRef,
+)
+from .field_walk import list_depth, terminal_of, walk_shape
+from .model_extraction import extract_model, resolve_field_alias
+from .specs import AnnotatedField, FieldSpec, MemberSpec, UnionSpec, is_model_class
+from .type_analyzer import (
+    capture_union_members,
+    single_literal_value,
+)
+
+__all__ = ["extract_discriminator", "extract_union"]
+
+
+def _find_common_base(members: list[type[BaseModel]]) -> type[BaseModel]:
+    """Find the most-derived common BaseModel ancestor of all members."""
+    if not members:
+        raise ValueError("Cannot find common base of empty members list")
+    filtered_mros = [
+        [c for c in cls.__mro__ if is_model_class(c) and c is not BaseModel]
+        for cls in members
+    ]
+    common = set(filtered_mros[0])
+    for mro in filtered_mros[1:]:
+        common &= set(mro)
+    if not common:
+        raise ValueError(
+            f"No common BaseModel ancestor for {[m.__name__ for m in members]}"
+        )
+
+    def max_mro_index(cls: type) -> int:
+        return max(mro.index(cls) for mro in filtered_mros)
+
+    return min(common, key=lambda c: (max_mro_index(c), c.__module__, c.__qualname__))
+
+
+def _find_field_by_alias(model: type[BaseModel], alias: str) -> FieldInfo | None:
+    """Find a field in `model_fields` by alias-resolved name."""
+    direct = model.model_fields.get(alias)
+    if direct is not None:
+        return direct
+    for py_name, fi in model.model_fields.items():
+        if resolve_field_alias(py_name, fi) == alias:
+            return fi
+    return None
+
+
+def extract_discriminator(
+    annotation: object,
+    members: list[type[BaseModel]],
+) -> tuple[str | None, dict[str, type[BaseModel]] | None]:
+    """Extract discriminator field name and value-to-type mapping."""
+    if get_origin(annotation) is not Annotated:
+        return None, None
+
+    disc_field_name: str | None = None
+    for metadata in get_args(annotation)[1:]:
+        if isinstance(metadata, FieldInfo):
+            disc_field_name = resolve_discriminator_field_name(metadata.discriminator)
+            if disc_field_name is not None:
+                break
+
+    if disc_field_name is None:
+        return None, None
+
+    mapping: dict[str, type[BaseModel]] = {}
+    for member in members:
+        field_info = _find_field_by_alias(member, disc_field_name)
+        if field_info and field_info.annotation is not None:
+            lit_val = single_literal_value(field_info.annotation)
+            if lit_val is not None:
+                key = lit_val.value if isinstance(lit_val, Enum) else str(lit_val)
+                mapping[key] = member
+
+    return disc_field_name, mapping or None
+
+
+_TypeShape = tuple[object, ...]
+_FieldKey = tuple[str, _TypeShape, frozenset[object]]
+
+
+def _structural_fingerprint(spec: FieldSpec) -> _TypeShape:
+    """Structural shape for dedup: ignores per-variant source_type variation.
+
+    Two fields with the same name and same `(terminal_base_type,
+    terminal_kind, is_optional, list_depth)` collapse to a single
+    `AnnotatedField` whose `variant_sources` lists the contributing
+    members.
+
+    `terminal_of` unwraps `ArrayOf` / `NewTypeShape`, so the terminal is
+    always one of the six leaf variants below; an unrecognized one
+    raises instead of silently collapsing into a shared fingerprint.
+    """
+    depth = list_depth(spec.shape)
+    base_type: object
+    terminal = terminal_of(spec.shape)
+    match terminal:
+        case Primitive(base_type=bt):
+            base_type, kind = bt, "scalar"
+        case LiteralScalar(values=values):
+            base_type, kind = ("Literal", values), "scalar"
+        case AnyScalar():
+            base_type, kind = "Any", "scalar"
+        case ModelRef(model=model):
+            base_type, kind = model.name, "model"
+        case UnionRef(union=union):
+            base_type, kind = union.name, "union"
+        case MapOf():
+            base_type, kind = "dict", "map"
+        case _:
+            raise TypeError(f"Unexpected terminal shape: {terminal!r}")
+    return (base_type, kind, spec.is_optional, depth)
+
+
+def _fingerprint_key(constraint: object) -> object:
+    """Return a value-stable set key for a single constraint.
+
+    Constraints with value equality -- every `FieldConstraint`, the
+    `annotated_types` dataclasses, `GeometryTypeConstraint` -- key as
+    themselves. Foreign metadata that falls back to identity equality, namely
+    pydantic's internal `Field(...)` metadata, keys on its value-stable `repr`
+    so two equal-valued instances still collapse.
+    """
+    if type(constraint).__eq__ is object.__eq__:
+        return repr(constraint)
+    return constraint
+
+
+def _constraints_fingerprint(spec: FieldSpec) -> frozenset[object]:
+    """Constraints declared anywhere in *spec*'s shape tree, as a comparable set.
+
+    `_structural_fingerprint` deliberately ignores constraints so that
+    members declaring the same field with per-variant `Annotated`
+    metadata still collapse to one `AnnotatedField`. This captures what
+    that ignores, so collisions with diverging constraints fail loudly
+    instead of silently keeping the last member's `FieldSpec`.
+
+    Constraint identity lives on the constraints themselves: `FieldConstraint`
+    subclasses define value equality and hashing, so equal rules collapse in
+    the set. `_fingerprint_key` covers the lone foreign holdout that still
+    compares by identity.
+    """
+    keys: list[object] = []
+
+    def collect(shape: FieldShape) -> None:
+        match shape:
+            case (
+                Primitive(constraints=cs)
+                | LiteralScalar(constraints=cs)
+                | AnyScalar(constraints=cs)
+                | ArrayOf(constraints=cs)
+                | MapOf(constraints=cs)
+            ):
+                for source in cs:
+                    keys.append(_fingerprint_key(source.constraint))
+            case ModelRef() | UnionRef() | NewTypeShape():
+                pass
+
+    walk_shape(spec.shape, collect)
+    return frozenset(keys)
+
+
+def extract_union(
+    name: str,
+    annotation: object,
+    *,
+    entry_point: str | None = None,
+    partitions: Mapping[str, str] | None = None,
+) -> UnionSpec:
+    """Extract a `UnionSpec` from a discriminated union type alias."""
+    extracted = capture_union_members(annotation)
+    if extracted is None:
+        raise TypeError(f"{name} is not a union type alias")
+    member_tuple, description = extracted
+    members = list(member_tuple)
+
+    common_base = _find_common_base(members)
+
+    # Plain Python type aliases (`Foo = Annotated[...]`) don't preserve
+    # the alias name in the annotation. The nested-union path (called
+    # from extract_model for UNION-kind fields) passes `members[0].__name__`
+    # as the placeholder name. Recover the alias by convention: members
+    # extend `<Alias>Base`, so stripping that suffix yields the alias.
+    # Top-level unions go through the CLI, which supplies the real name
+    # and skips this fallback.
+    #
+    # PEP 695 (`type Foo = Annotated[...]`) preserves `__name__` as
+    # `"Foo"` on 3.12+; after migrating, the placeholder hack can go.
+    member_names = {m.__name__ for m in members}
+    if name in member_names:
+        base_name = common_base.__name__
+        name = (
+            base_name.removesuffix("Base") if base_name.endswith("Base") else base_name
+        )
+
+    base_spec = extract_model(common_base)
+    shared_field_names = {f.name for f in base_spec.fields}
+
+    member_specs = [MemberSpec(m, extract_model(m)) for m in members]
+
+    annotated_fields: list[AnnotatedField] = []
+
+    for fs in base_spec.fields:
+        annotated_fields.append(AnnotatedField(field_spec=fs, variant_sources=None))
+
+    seen: dict[_FieldKey, AnnotatedField] = {}
+
+    for member in member_specs:
+        member_cls = member.member_cls
+        for fs in member.spec.fields:
+            if fs.name in shared_field_names:
+                continue
+            # The key includes the constraints fingerprint alongside the
+            # structural one: two arms with the same name and shape but
+            # different constraints (e.g. VehicleAxleCountSelector's
+            # `ge=1, le=100, multiple_of=1` vs the other selectors' `ge=0`)
+            # must not collapse into one `AnnotatedField` sharing a single
+            # constraint set -- that would silently drop one arm's rules.
+            # Keeping them as separate rows, each gated to its own
+            # `variant_sources`, reuses the same per-arm `Guard` mechanism
+            # that already handles a field present on only some arms
+            # (`check_builder._field_checks_for_union`), and the renderer's
+            # collision resolver already disambiguates multiple `Check`s
+            # landing on the same field label.
+            key = (fs.name, _structural_fingerprint(fs), _constraints_fingerprint(fs))
+            existing = seen.get(key)
+            prior_sources = existing.variant_sources or () if existing else ()
+            seen[key] = AnnotatedField(
+                field_spec=fs,
+                variant_sources=(*prior_sources, member_cls),
+            )
+
+    annotated_fields.extend(seen.values())
+
+    disc_field, disc_mapping = extract_discriminator(annotation, members)
+
+    return UnionSpec(
+        name=name,
+        description=description,
+        annotated_fields=annotated_fields,
+        members=members,
+        member_specs=member_specs,
+        discriminator_field=disc_field,
+        discriminator_mapping=disc_mapping,
+        source_annotation=annotation,
+        common_base=common_base,
+        entry_point=entry_point,
+        partitions=partitions or {},
+    )

overture/schema/codegen/layout/module_layout.py

@@ -0,0 +1,139 @@
+"""Output directory layout from Python module paths.
+
+Translates dotted module paths into output directory paths by mirroring
+the source package structure.
+"""
+
+from __future__ import annotations
+
+import sys
+from collections.abc import Iterable, Mapping
+from pathlib import PurePosixPath
+
+from overture.schema.system.discovery import split_entry_point
+
+__all__ = [
+    "OUTPUT_ROOT",
+    "compute_output_dir",
+    "compute_schema_root",
+    "entry_point_class",
+    "entry_point_module",
+    "is_package_module",
+    "module_relpath",
+    "output_dir_for_entry_point",
+]
+
+OUTPUT_ROOT = PurePosixPath(".")
+
+
+def entry_point_module(entry_point_path: str) -> str:
+    """Extract module path from entry-point-style path.
+
+    >>> entry_point_module("overture.schema.buildings:Building")
+    'overture.schema.buildings'
+    """
+    return split_entry_point(entry_point_path)[0]
+
+
+def entry_point_class(entry_point_path: str) -> str:
+    """Extract class name from entry-point-style path.
+
+    >>> entry_point_class("overture.schema.buildings:Building")
+    'Building'
+    """
+    return split_entry_point(entry_point_path)[1]
+
+
+def compute_schema_root(module_paths: Iterable[str]) -> str:
+    """Find the longest common dotted prefix of module paths.
+
+    Deduplicates inputs first. For a single unique path, drops the
+    last component (the module itself).
+    """
+    paths = sorted(set(module_paths))
+    if not paths:
+        msg = "No module paths provided"
+        raise ValueError(msg)
+
+    segments = [p.split(".") for p in paths]
+    if len(segments) == 1:
+        return ".".join(segments[0][:-1])
+
+    common: list[str] = []
+    for parts in zip(*segments, strict=False):
+        if len(set(parts)) == 1:
+            common.append(parts[0])
+        else:
+            break
+    return ".".join(common)
+
+
+def module_relpath(module: str, root: str) -> str:
+    """Strip the schema root prefix from a dotted module path."""
+    if not root:
+        return module
+    if module == root:
+        return ""
+    prefix = root + "."
+    if not module.startswith(prefix):
+        msg = f"Module {module!r} does not start with root {root!r}"
+        raise ValueError(msg)
+    return module[len(prefix) :]
+
+
+def is_package_module(
+    module: str,
+    module_registry: Mapping[str, object] | None = None,
+) -> bool:
+    """Check whether a module is a package (directory) or a file module.
+
+    Packages have `__path__`; file modules do not (PEP 302).
+    """
+    registry: Mapping[str, object] = (
+        module_registry if module_registry is not None else sys.modules
+    )
+    mod = registry.get(module)
+    if mod is None:
+        msg = f"Module {module!r} not found in registry"
+        raise ValueError(msg)
+    return hasattr(mod, "__path__")
+
+
+def output_dir_for_entry_point(
+    entry_point_path: str | None,
+    schema_root: str,
+    module_registry: Mapping[str, object] | None = None,
+) -> PurePosixPath:
+    """Compute output directory from an entry-point-style path.
+
+    Raises ValueError if *entry_point_path* is None.
+    """
+    if entry_point_path is None:
+        msg = "entry_point_path must not be None"
+        raise ValueError(msg)
+    module = entry_point_module(entry_point_path)
+    return compute_output_dir(module, schema_root, module_registry)
+
+
+def compute_output_dir(
+    module: str,
+    schema_root: str,
+    module_registry: Mapping[str, object] | None = None,
+) -> PurePosixPath:
+    """Compute output directory for a module, mirroring package structure.
+
+    File modules drop their last component (the .py filename).
+    Packages keep all components. Returns `PurePosixPath(".")` for
+    the root directory.
+    """
+    relpath = module_relpath(module, schema_root)
+    if not relpath:
+        return OUTPUT_ROOT
+
+    parts = relpath.split(".")
+    if not is_package_module(module, module_registry):
+        parts = parts[:-1]
+
+    if not parts:
+        return OUTPUT_ROOT
+    return PurePosixPath(*parts)

overture/schema/codegen/layout/type_collection.py

@@ -0,0 +1,122 @@
+"""Supplementary type discovery by walking feature trees.
+
+Walks `FieldShape` trees to extract referenced enums, NewTypes,
+Pydantic built-ins, and union member sub-models. `ModelRef` and
+`UnionRef` carry their resolved specs structurally, so recursion
+follows the shape directly.
+"""
+
+from collections.abc import Sequence
+from enum import Enum
+from typing import Annotated, get_args, get_origin
+
+from pydantic import BaseModel
+
+from ..extraction.enum_extraction import extract_enum
+from ..extraction.field import (
+    FieldShape,
+    ModelRef,
+    NewTypeShape,
+    Primitive,
+    UnionRef,
+)
+from ..extraction.field_walk import walk_shape
+from ..extraction.newtype_extraction import extract_newtype
+from ..extraction.pydantic_extraction import extract_pydantic_type
+from ..extraction.specs import (
+    FieldSpec,
+    ModelSpec,
+    RecordSpec,
+    SupplementarySpec,
+    TypeIdentity,
+    is_pydantic_sourced,
+)
+from ..extraction.type_analyzer import analyze_type, is_newtype
+from ..extraction.type_registry import is_semantic_newtype
+
+__all__ = ["collect_all_supplementary_types"]
+
+
+def collect_all_supplementary_types(
+    model_specs: Sequence[ModelSpec],
+) -> dict[TypeIdentity, SupplementarySpec]:
+    """Collect supplementary types by walking expanded feature trees.
+
+    Walks `ModelRef` references for sub-models (already extracted),
+    and extracts enums and NewTypes on first encounter. Two types
+    with the same class name from different modules are keyed
+    separately.
+    """
+    feature_objs: set[object] = {spec.identity.obj for spec in model_specs}
+    all_specs: dict[TypeIdentity, SupplementarySpec] = {}
+    visited_models: set[object] = set()
+
+    def _register_newtype(newtype_ref: object, name: str) -> bool:
+        nt_id = TypeIdentity(newtype_ref, name)
+        if nt_id in all_specs:
+            return False
+        all_specs[nt_id] = extract_newtype(newtype_ref)
+        return True
+
+    def _collect_from_model(model_spec: RecordSpec) -> None:
+        if (
+            model_spec.source_type in visited_models
+            or model_spec.source_type in feature_objs
+        ):
+            return
+        visited_models.add(model_spec.source_type)
+        all_specs[model_spec.identity] = model_spec
+        _collect_from_fields(model_spec.fields)
+
+    def _collect_inner_newtypes(newtype_ref: object) -> None:
+        """Walk a NewType's `__supertype__` chain for nested semantic NewTypes."""
+        annotation = getattr(newtype_ref, "__supertype__", None)
+        while annotation is not None:
+            if get_origin(annotation) is Annotated:
+                annotation = get_args(annotation)[0]
+                continue
+            if is_newtype(annotation):
+                inner_shape, _, _ = analyze_type(annotation)
+                if isinstance(inner_shape, NewTypeShape) and is_semantic_newtype(
+                    inner_shape
+                ):
+                    _register_newtype(inner_shape.ref, inner_shape.name)
+                annotation = getattr(annotation, "__supertype__", None)
+                continue
+            break
+
+    def _collect_from_shape(shape: FieldShape) -> None:
+        """Walk *shape* and register every supplementary type it touches."""
+
+        def _visit(node: FieldShape) -> None:
+            match node:
+                case NewTypeShape(name=name, ref=ref) if is_semantic_newtype(node):
+                    if _register_newtype(ref, name):
+                        _collect_inner_newtypes(ref)
+                case UnionRef(union=u):
+                    for member in u.member_specs:
+                        _collect_from_model(member.spec)
+                case ModelRef(model=m, starts_cycle=False):
+                    _collect_from_model(m)
+                case Primitive(source_type=cls) if cls is not None and isinstance(
+                    cls, type
+                ):
+                    if issubclass(cls, Enum):
+                        eid = TypeIdentity.of(cls)
+                        if eid not in all_specs:
+                            all_specs[eid] = extract_enum(cls)
+                    elif is_pydantic_sourced(cls) and not issubclass(cls, BaseModel):
+                        pid = TypeIdentity.of(cls)
+                        if pid not in all_specs:
+                            all_specs[pid] = extract_pydantic_type(cls)
+
+        walk_shape(shape, _visit)
+
+    def _collect_from_fields(fields: list[FieldSpec]) -> None:
+        for field_spec in fields:
+            _collect_from_shape(field_spec.shape)
+
+    for spec in model_specs:
+        _collect_from_fields(spec.fields)
+
+    return all_specs