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

overture/schema/codegen/extraction/literal_alternatives.py

@@ -0,0 +1,26 @@
+"""Internal constraint recording a union's literal alternatives.
+
+A field annotated `X | Literal[c, ...]` validates as "the concrete arm `X`'s
+checks pass OR the value is one of `c, ...`". `type_analyzer._peel_union` keeps
+the concrete arm as the field's shape (so downstream consumers still see a
+`Primitive` / `NewTypeShape` rather than a union of scalar-and-literal) and
+records the dropped literal values in this constraint on that shape's layer.
+
+Consumers read it to let those literal values bypass the concrete arm's
+constraints: the PySpark dispatch emits a value-exact bypass, and the markdown
+renderer notes the accepted literals. Codegen-internal -- schema authors write
+the plain `X | Literal[c]` union; nothing constructs this class directly.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+__all__ = ["LiteralAlternatives"]
+
+
+@dataclass(frozen=True, slots=True)
+class LiteralAlternatives:
+    """Literal values a union field accepts alongside its concrete arm."""
+
+    values: tuple[object, ...]

overture/schema/codegen/extraction/model_constraints.py

@@ -0,0 +1,252 @@
+"""Convert model-level constraints to human-readable prose.
+
+Handles RequireAnyOf, RadioGroup, ForbidIf, RequireIf, and other
+ModelConstraint types. Produces descriptions and per-field notes for
+documentation rendering.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from overture.schema.system.model_constraint import (
+    FieldEqCondition,
+    ForbidIfConstraint,
+    MinFieldsSetConstraint,
+    ModelConstraint,
+    NoExtraFieldsConstraint,
+    Not,
+    RadioGroupConstraint,
+    RequireAnyOfConstraint,
+    RequireAnyTrueConstraint,
+    RequireIfConstraint,
+)
+
+__all__ = ["analyze_model_constraints"]
+
+_ConditionalConstraint = RequireIfConstraint | ForbidIfConstraint
+
+
+@dataclass(frozen=True)
+class _ConstraintEntry:
+    """A constraint description paired with the field names it affects."""
+
+    description: str
+    field_names: frozenset[str]
+
+
+def _backtick_join(names: tuple[str, ...]) -> str:
+    """Format field names as backtick-quoted, comma-separated list."""
+    return ", ".join(f"`{n}`" for n in names)
+
+
+def _conditional_verb(constraint: _ConditionalConstraint) -> str:
+    """Return 'required' or 'forbidden' based on constraint type."""
+    return "required" if isinstance(constraint, RequireIfConstraint) else "forbidden"
+
+
+def _plural_verb(names: tuple[str, ...]) -> str:
+    """Return 'is' or 'are' based on field count."""
+    return "are" if len(names) > 1 else "is"
+
+
+def _unwrap_field_eq(condition: object) -> tuple[FieldEqCondition, bool] | None:
+    """Extract the FieldEqCondition from a condition, with negation flag.
+
+    Returns (field_eq, is_negated) or None for unrecognized conditions.
+    """
+    if isinstance(condition, Not) and isinstance(condition.inner, FieldEqCondition):
+        return condition.inner, True
+    if isinstance(condition, FieldEqCondition):
+        return condition, False
+    return None
+
+
+def _describe_condition(condition: object) -> str:
+    """Render a Condition as human-readable text."""
+    unwrapped = _unwrap_field_eq(condition)
+    if unwrapped is not None:
+        field_eq, negated = unwrapped
+        op = "≠" if negated else "="
+        return f"`{field_eq.field_name}` {op} `{field_eq.value}`"
+    return str(condition)
+
+
+def _unwrap_true_field_eq(condition: object) -> FieldEqCondition | None:
+    if isinstance(condition, FieldEqCondition) and condition.value is True:
+        return condition
+    return None
+
+
+def _describe_conditional(constraint: _ConditionalConstraint) -> str:
+    """Describe a require_if or forbid_if constraint."""
+    fields = _backtick_join(constraint.field_names)
+    verb = _conditional_verb(constraint)
+    cond = _describe_condition(constraint.condition)
+    return f"{fields} {_plural_verb(constraint.field_names)} {verb} when {cond}"
+
+
+def _consolidation_key(
+    constraint: _ConditionalConstraint,
+) -> tuple[type, tuple[str, ...], str] | None:
+    """Return a grouping key if the constraint is consolidatable, else None.
+
+    Consolidatable: same type, same field_names, plain FieldEqCondition
+    (not negated) on the same condition field.
+    """
+    cond = constraint.condition
+    if not isinstance(cond, FieldEqCondition):
+        return None
+    return (type(constraint), constraint.field_names, cond.field_name)
+
+
+def _as_field_eq(constraint: _ConditionalConstraint) -> FieldEqCondition:
+    """Narrow a conditional constraint's condition to FieldEqCondition.
+
+    Only called on constraints that passed _consolidation_key, which
+    rejects non-FieldEqCondition conditions.
+    """
+    cond = constraint.condition
+    if not isinstance(cond, FieldEqCondition):
+        raise TypeError(f"Expected FieldEqCondition, got {type(cond).__name__}")
+    return cond
+
+
+def _describe_consolidated(
+    constraints: list[_ConditionalConstraint],
+) -> str:
+    """Describe a group of consolidated conditional constraints."""
+    first = constraints[0]
+    fields = _backtick_join(first.field_names)
+    verb = _conditional_verb(first)
+    cond_field = _as_field_eq(first).field_name
+    values = ", ".join(f"`{_as_field_eq(c).value}`" for c in constraints)
+    return (
+        f"{fields} {_plural_verb(first.field_names)} {verb} "
+        f"when `{cond_field}` is one of: {values}"
+    )
+
+
+def _condition_field_names(condition: object) -> frozenset[str]:
+    """Extract field names referenced by a condition."""
+    unwrapped = _unwrap_field_eq(condition)
+    if unwrapped is not None:
+        return frozenset({unwrapped[0].field_name})
+    return frozenset()
+
+
+def _affected_field_names(constraint: ModelConstraint) -> frozenset[str]:
+    """Return all field names referenced by a constraint.
+
+    Includes both constrained field_names and condition trigger fields.
+    Returns empty set for constraints that don't reference specific fields
+    (NoExtraFieldsConstraint, MinFieldsSetConstraint).
+    """
+    if isinstance(constraint, (NoExtraFieldsConstraint, MinFieldsSetConstraint)):
+        return frozenset()
+    if isinstance(constraint, (RequireIfConstraint, ForbidIfConstraint)):
+        return frozenset(constraint.field_names) | _condition_field_names(
+            constraint.condition
+        )
+    if isinstance(constraint, (RequireAnyOfConstraint, RadioGroupConstraint)):
+        return frozenset(constraint.field_names)
+    if isinstance(constraint, RequireAnyTrueConstraint):
+        return frozenset().union(
+            *(_condition_field_names(condition) for condition in constraint.conditions)
+        )
+    return frozenset()
+
+
+def _describe_one(constraint: ModelConstraint) -> str | None:
+    """Describe a single constraint, or None to skip it."""
+    if isinstance(constraint, NoExtraFieldsConstraint):
+        return None
+    if isinstance(constraint, RequireAnyOfConstraint):
+        return f"At least one of {_backtick_join(constraint.field_names)} must be set"
+    if isinstance(constraint, RequireAnyTrueConstraint):
+        true_field_conditions = [
+            field_eq
+            for condition in constraint.conditions
+            if (field_eq := _unwrap_true_field_eq(condition)) is not None
+        ]
+        if len(true_field_conditions) == len(constraint.conditions):
+            return (
+                "At least one of "
+                f"{_backtick_join(tuple(c.field_name for c in true_field_conditions))} "
+                "must be `true`"
+            )
+        return "At least one of these conditions must be true: " + ", ".join(
+            _describe_condition(condition) for condition in constraint.conditions
+        )
+    if isinstance(constraint, RadioGroupConstraint):
+        return f"Exactly one of {_backtick_join(constraint.field_names)} must be `true`"
+    if isinstance(constraint, MinFieldsSetConstraint):
+        return f"At least {constraint.count} fields must be set"
+    if isinstance(constraint, (RequireIfConstraint, ForbidIfConstraint)):
+        return _describe_conditional(constraint)
+    return f"`{constraint.name}`"
+
+
+def _analyze_constraints(
+    constraints: tuple[ModelConstraint, ...],
+) -> list[_ConstraintEntry]:
+    """Analyze constraints into descriptions paired with affected fields.
+
+    Handles consolidation and filtering, preserving original declaration order.
+    """
+    groups: dict[
+        tuple[type, tuple[str, ...], str], list[tuple[int, _ConditionalConstraint]]
+    ] = {}
+    standalone: list[tuple[int, ModelConstraint]] = []
+
+    for i, c in enumerate(constraints):
+        if isinstance(c, (RequireIfConstraint, ForbidIfConstraint)):
+            key = _consolidation_key(c)
+            if key is not None:
+                groups.setdefault(key, []).append((i, c))
+                continue
+        standalone.append((i, c))
+
+    entries: list[tuple[int, _ConstraintEntry]] = []
+
+    for group_items in groups.values():
+        first_idx = group_items[0][0]
+        group_constraints = [c for _, c in group_items]
+        all_fields: frozenset[str] = frozenset().union(
+            *(_affected_field_names(c) for c in group_constraints)
+        )
+        if len(group_constraints) == 1:
+            desc = _describe_one(group_constraints[0])
+        else:
+            desc = _describe_consolidated(group_constraints)
+        if desc is not None:
+            entries.append((first_idx, _ConstraintEntry(desc, all_fields)))
+
+    for idx, c in standalone:
+        desc = _describe_one(c)
+        if desc is not None:
+            entries.append((idx, _ConstraintEntry(desc, _affected_field_names(c))))
+
+    entries.sort(key=lambda e: e[0])
+    return [entry for _, entry in entries]
+
+
+def analyze_model_constraints(
+    constraints: tuple[ModelConstraint, ...],
+) -> tuple[list[str], dict[str, list[str]]]:
+    """Analyze constraints into descriptions and per-field notes in one pass.
+
+    Returns (descriptions, field_notes) where descriptions is the list of
+    human-readable constraint strings and field_notes maps field names to
+    constraint descriptions that reference them.
+    """
+    entries = _analyze_constraints(constraints)
+
+    descriptions = [entry.description for entry in entries]
+
+    field_notes: dict[str, list[str]] = {}
+    for entry in entries:
+        for name in entry.field_names:
+            field_notes.setdefault(name, []).append(entry.description)
+
+    return descriptions, field_notes

overture/schema/codegen/extraction/model_extraction.py

@@ -0,0 +1,240 @@
+"""Pydantic model extraction into `RecordSpec`."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+from pydantic import BaseModel
+from pydantic.fields import FieldInfo
+from pydantic_core import PydanticUndefined
+
+from overture.schema.system.model_constraint import ModelConstraint
+
+from .docstring import clean_docstring
+from .field import (
+    ModelRef,
+    UnionRef,
+)
+from .specs import FieldSpec, RecordSpec, is_model_class
+from .type_analyzer import (
+    ModelResolver,
+    UnionResolver,
+    analyze_type,
+    attach_field_metadata,
+    unwrap_list,
+)
+
+__all__ = [
+    "extract_model",
+    "resolve_field_alias",
+]
+
+
+def resolve_field_alias(field_name: str, field_info: FieldInfo) -> str:
+    """Return the data-dict key for a Pydantic field.
+
+    Prefers `validation_alias`, falls back to `alias`, then the
+    Python field name. Only string aliases are supported; AliasPath
+    and AliasChoices are ignored.
+    """
+    validation_alias = field_info.validation_alias
+    if isinstance(validation_alias, str):
+        return validation_alias
+    alias = field_info.alias
+    if isinstance(alias, str):
+        return alias
+    return field_name
+
+
+def _is_field_required(field_info: FieldInfo, is_optional: bool) -> bool:
+    """Determine whether a field is required (no default and not Optional)."""
+    has_default = (
+        field_info.default is not PydanticUndefined
+        or field_info.default_factory is not None
+    )
+    return not has_default and not is_optional
+
+
+def _basemodel_bases(cls: type) -> list[type[BaseModel]]:
+    """Return direct BaseModel bases, excluding BaseModel itself."""
+    return [b for b in cls.__bases__ if is_model_class(b) and b is not BaseModel]
+
+
+def _class_order(model_class: type[BaseModel]) -> list[type]:
+    """Return MRO classes in documentation order, recursively.
+
+    For single-inheritance: reversed MRO (base first, derived last).
+    For multiple-inheritance: primary chain → self → mixins, where
+    primary chain and each mixin are themselves recursively ordered.
+    """
+    bases = _basemodel_bases(model_class)
+
+    if len(bases) <= 1:
+        return [
+            cls
+            for cls in reversed(model_class.__mro__)
+            if issubclass(cls, BaseModel) and cls is not BaseModel
+        ]
+
+    primary = _class_order(bases[0])
+    mixins = [cls for base in bases[1:] for cls in _class_order(base)]
+    return primary + [model_class] + mixins
+
+
+def _field_order(model_class: type[BaseModel]) -> list[str]:
+    """Return `model_fields` keys in documentation order.
+
+    Walks the class hierarchy recursively. At each level of multiple
+    inheritance, the first base is the primary chain and the rest are
+    mixins. Primary chain and own fields come first, then mixin fields
+    in declaration order. Single-inheritance levels use Pydantic's
+    default reversed-MRO order.
+    """
+    valid_names = set(model_class.model_fields.keys())
+    result: list[str] = []
+    seen: set[str] = set()
+    for cls in _class_order(model_class):
+        for name in getattr(cls, "__annotations__", {}):
+            if name not in seen and name in valid_names:
+                result.append(name)
+                seen.add(name)
+    return result
+
+
+def extract_model(
+    model_class: type[BaseModel],
+    *,
+    entry_point: str | None = None,
+    partitions: Mapping[str, str] | None = None,
+) -> RecordSpec:
+    """Extract a fully-resolved `RecordSpec` from a Pydantic model class.
+
+    Recurses into sub-models and unions, producing `ModelRef` /
+    `UnionRef` terminals with their specs resolved. Cycles in the
+    model graph (a field whose source type is an ancestor on the
+    current extraction stack) produce a `ModelRef` pointing at the
+    in-progress ancestor spec with `starts_cycle=True` so consumers
+    stop recursion at the back-edge.
+
+    The caller must not pass a `RootModel`: it serializes as its bare
+    root value and has no record structure, so walking its fields yields
+    only a spurious `root` column. `extract_model_spec` filters RootModel
+    entry points before they reach here, and a RootModel reached as a
+    field is unwrapped to its bare shape by `analyze_type`.
+    """
+    return _extract_model_recursive(
+        model_class,
+        entry_point=entry_point,
+        partitions=partitions or {},
+        cache={},
+        ancestors=frozenset(),
+    )
+
+
+def _extract_model_recursive(
+    model_class: type[BaseModel],
+    *,
+    entry_point: str | None,
+    partitions: Mapping[str, str],
+    cache: dict[type, RecordSpec],
+    ancestors: frozenset[type],
+) -> RecordSpec:
+    """Inner recursive helper for `extract_model`.
+
+    Inserts the (partial) `RecordSpec` into `cache` before populating
+    its fields so cycles can find it. `ancestors` is the set of types
+    currently on the recursion stack -- a sub-field whose source type
+    appears there is a back-edge and gets `starts_cycle=True`.
+    """
+    spec = RecordSpec(
+        name=model_class.__name__,
+        description=clean_docstring(model_class.__doc__),
+        fields=[],
+        source_type=model_class,
+        entry_point=entry_point,
+        partitions=partitions,
+        constraints=ModelConstraint.get_model_constraints(model_class),
+    )
+    cache[model_class] = spec
+    descendant_ancestors = ancestors | {model_class}
+
+    model_resolver, union_resolver = _make_resolvers(cache, descendant_ancestors)
+
+    fields: list[FieldSpec] = []
+    for field_name in _field_order(model_class):
+        field_info = model_class.model_fields[field_name]
+        annotation = field_info.annotation
+        if annotation is None:
+            continue
+        shape, is_optional, ti_description = analyze_type(
+            annotation,
+            owner=model_class,
+            model_resolver=model_resolver,
+            union_resolver=union_resolver,
+        )
+        # Pydantic strips the outermost Annotated wrapper from some fields
+        # (non-optional, non-union) and moves its metadata to
+        # `field_info.metadata`; `analyze_type` then sees a bare type and
+        # misses those constraints. Reattach them at the topmost
+        # constraint-bearing layer.
+        shape = attach_field_metadata(shape, field_info)
+        fields.append(
+            FieldSpec(
+                name=resolve_field_alias(field_name, field_info),
+                shape=shape,
+                description=field_info.description or ti_description,
+                is_required=_is_field_required(field_info, is_optional),
+                is_optional=is_optional,
+            )
+        )
+
+    spec.fields = fields
+    return spec
+
+
+def _make_resolvers(
+    cache: dict[type, RecordSpec],
+    ancestors: frozenset[type],
+) -> tuple[ModelResolver, UnionResolver]:
+    """Build the resolvers that recursively extract sub-models / sub-unions.
+
+    `cache` shares already-extracted sub-specs across a single
+    extraction so sub-models referenced more than once share a
+    `RecordSpec`. `ancestors` carries the recursion stack for cycle
+    detection -- a back-edge produces a `ModelRef` pointing at the
+    in-progress ancestor spec with `starts_cycle=True`.
+    """
+
+    def resolve_model(cls: type[BaseModel]) -> ModelRef:
+        if cls in ancestors:
+            return ModelRef(model=cache[cls], starts_cycle=True)
+        cached = cache.get(cls)
+        if cached is not None:
+            return ModelRef(model=cached)
+        sub_spec = _extract_model_recursive(
+            cls,
+            entry_point=None,
+            partitions={},
+            cache=cache,
+            ancestors=ancestors,
+        )
+        return ModelRef(model=sub_spec)
+
+    def resolve_union(
+        annotation: object,
+        members: tuple[type[BaseModel], ...],
+        _description: str | None,
+    ) -> UnionRef:
+        # Late import: extract_union calls back into extract_model for
+        # member classes. A module-level import would be a cycle.
+        from .union_extraction import extract_union  # noqa: PLC0415
+
+        # Recover the union alias name: `analyze_type` reaches the
+        # union via `members[0].__name__` when the alias name is lost
+        # (plain `Foo = Annotated[...]` doesn't preserve it pre-PEP-695).
+        # Convention: members extend `<Alias>Base`.
+        placeholder = members[0].__name__ if members else ""
+        sub_union = extract_union(placeholder, unwrap_list(annotation))
+        return UnionRef(union=sub_union)
+
+    return resolve_model, resolve_union

overture/schema/codegen/extraction/newtype_extraction.py

@@ -0,0 +1,73 @@
+"""Type-alias extraction: NewType and RootModel into `NewTypeSpec`.
+
+Both a NewType and a `RootModel` are named aliases over an underlying
+type -- a RootModel serializes as its bare root value -- so both document
+as the same `NewTypeSpec` (a name plus the underlying shape).
+"""
+
+from pydantic import RootModel
+
+from .docstring import clean_docstring, is_custom_docstring
+from .field import NewTypeShape
+from .specs import NewTypeSpec
+from .type_analyzer import analyze_type
+
+__all__ = ["extract_newtype", "extract_rootmodel_alias"]
+
+
+def extract_newtype(newtype_callable: object) -> NewTypeSpec:
+    """Extract a `NewTypeSpec` from a NewType callable.
+
+    `analyze_type(newtype_callable)` returns a shape whose outermost
+    layer is the NewType's own `NewTypeShape`. We strip that wrapper so
+    `NewTypeSpec.shape` describes the *underlying* type -- the NewType
+    isn't a self-reference on its own page.
+    """
+    shape, _, ti_description = analyze_type(newtype_callable)
+
+    name = getattr(newtype_callable, "__name__", None)
+    if isinstance(shape, NewTypeShape) and shape.name == name:
+        underlying = shape.inner
+    else:
+        underlying = shape
+
+    if name is None:
+        msg = f"Cannot determine name for NewType: {newtype_callable!r}"
+        raise ValueError(msg)
+
+    doc = getattr(newtype_callable, "__doc__", None)
+    description = clean_docstring(doc) if is_custom_docstring(doc) else ti_description
+
+    return NewTypeSpec(
+        name=name,
+        description=description,
+        shape=underlying,
+        source_type=newtype_callable,
+    )
+
+
+def extract_rootmodel_alias(root_model: type[RootModel]) -> NewTypeSpec:
+    """Extract a `NewTypeSpec` documenting a `RootModel` as a named alias.
+
+    `analyze_type` already unwraps a RootModel to its bare root shape, so
+    there is no wrapper to strip -- unlike `extract_newtype`. A custom
+    class docstring names the alias; otherwise the root field's own
+    description stands in. Pydantic moves that description onto the root
+    `FieldInfo` (as it does for model fields), so read it there rather than
+    from `analyze_type`, which recurses past it.
+
+    No `is_custom_docstring` guard is needed (again unlike `extract_newtype`):
+    a RootModel subclass with no docstring has `__doc__ = None` -- classes do
+    not inherit `RootModel`'s base docstring -- so there is no auto-generated
+    text to filter out, and `clean_docstring(None)` falls through to the root
+    description.
+    """
+    shape, _, _ = analyze_type(root_model)
+    root = root_model.model_fields["root"]
+    description = clean_docstring(root_model.__doc__) or root.description
+    return NewTypeSpec(
+        name=root_model.__name__,
+        description=description,
+        shape=shape,
+        source_type=root_model,
+    )

overture/schema/codegen/extraction/numeric_extraction.py

@@ -0,0 +1,74 @@
+"""Numeric type extraction."""
+
+from annotated_types import Interval
+
+from .docstring import first_docstring_line
+from .field import FieldShape, Scalar
+from .field_walk import terminal_of
+from .newtype_extraction import extract_newtype
+from .specs import NumericSpec, TypeIdentity
+
+__all__ = [
+    "extract_numeric_bounds",
+    "extract_numerics",
+]
+
+
+# Bound attribute names on annotated_types constraints (Ge, Gt, Le, Lt, Interval).
+_BOUND_ATTRS = ("ge", "gt", "le", "lt")
+
+
+def extract_numeric_bounds(shape: FieldShape) -> Interval:
+    """Extract numeric bounds from the constraints on a shape's terminal scalar.
+
+    Walks `NewTypeShape` / `ArrayOf` wrappers to find the terminal
+    `Scalar`, then scans its constraints for `ge`, `gt`, `le`, and `lt`
+    attributes. Stops at the first constraint defining each bound.
+    """
+    terminal = terminal_of(shape)
+    if not isinstance(terminal, Scalar):
+        return Interval()
+    found: dict[str, int | float] = {}
+    for cs in terminal.constraints:
+        c = cs.constraint
+        for attr in _BOUND_ATTRS:
+            if attr not in found:
+                val = getattr(c, attr, None)
+                if val is not None:
+                    found[attr] = val
+    return Interval(**found)
+
+
+def extract_numerics(
+    numeric_ids: list[TypeIdentity],
+) -> list[NumericSpec]:
+    """Extract specifications for numeric types."""
+    specs: list[NumericSpec] = []
+    for tid in numeric_ids:
+        newtype_spec = extract_newtype(tid.obj)
+        # extract_newtype strips the outer NewTypeShape, so the spec's
+        # terminal scalar already carries the constraints the NewType
+        # contributed -- extract_numeric_bounds walks straight to it.
+        bounds = extract_numeric_bounds(newtype_spec.shape)
+        description = first_docstring_line(getattr(tid.obj, "__doc__", None))
+        float_bits = _extract_float_bits(tid.name)
+        specs.append(
+            NumericSpec(
+                name=tid.name,
+                description=description,
+                bounds=bounds,
+                float_bits=float_bits,
+            )
+        )
+    return specs
+
+
+_FLOAT_BITS: dict[str, int] = {
+    "float32": 32,
+    "float64": 64,
+}
+
+
+def _extract_float_bits(name: str) -> int | None:
+    """Extract bit width from a float type name like `float32`."""
+    return _FLOAT_BITS.get(name)