Schema Models¶

class SchemaModel(BaseModel):
    """Top-level schema definition."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    @model_validator(mode="before")
    @classmethod
    def reject_settings_moved_to_config(cls, data: Any) -> Any:
        """Point a schema that still declares a run setting at `nyctea.Config`.

        Args:
            data: Raw input, before field validation.

        Returns:
            The input unchanged.

        Raises:
            ValueError: If the input declares a setting that moved to `nyctea.Config`.
        """
        if isinstance(data, Mapping):
            moved = [name for name in _MOVED_TO_CONFIG if name in data]
            if moved:
                names = ", ".join(f"`{name}`" for name in moved)
                calls = ", ".join(f"nyctea.Config.set_{name}(...)" for name in moved)
                scoped = ", ".join(f"{name}=..." for name in moved)
                raise ValueError(
                    f"{names} moved from the schema to `nyctea.Config`. These describe the run "
                    f"rather than what valid data looks like. Remove them from the schema and use "
                    f"{calls}, or a `with nyctea.Config({scoped})` block."
                )
        return data

    column_matching: Literal["exact", "cleaned"] = Field(
        "exact",
        description=(
            "How a frame's column names are matched against canonical names and "
            "synonyms. 'exact' is the default and is unchanged. 'cleaned' also "
            "accepts names that match after trimming whitespace and Unicode case "
            "folding, so ' AGE ' resolves to 'age' without enumerating it as a "
            "synonym. Exact matches always win over cleaned ones, and two physical "
            "columns matching one schema column is still ambiguous rather than a "
            "guess."
        ),
    )

    coerce: bool = Field(
        True,
        description="Whether to coerce columns to the specified dtypes after parsing and validation",
    )

    on_failure: OnFailureBehavior = Field(
        "raise",
        description=(
            "Default failure handling for all columns:\n"
            "- 'raise': error, stop\n"
            "- 'null': failure value becomes or remains null\n"
            "- 'ignore': parser/coercion nulls remain, check failures are kept and reported"
        ),
    )

    columns: dict[str, ColumnSchema] = Field(..., description="Mapping of column name to its validation schema")
    frame_parsers: list[FrameParser] = Field(default_factory=list, description="DataFrame-level parsing functions")

    frame_checks: list[FrameCheck] = Field(default_factory=list, description="DataFrame-level checks")

    @model_validator(mode="after")
    def verify_name_ownership(self) -> "SchemaModel":
        """Ensure every accepted column name is claimed by exactly one column.

        A schema whose names overlap cannot be resolved unambiguously, so it is
        rejected here rather than at validation time. Reports every conflict at
        once, so a schema with several does not take several rounds to fix.

        Raises:
            ValueError: If any accepted name is claimed more than once.
        """
        claims: dict[str, list[str]] = {}
        for canonical, column in self.columns.items():
            claims.setdefault(canonical, []).append(canonical)
            for synonym in column.synonyms:
                claims.setdefault(synonym, []).append(canonical)

        conflicts: list[str] = []
        for name, claimants in claims.items():
            if len(claimants) == 1:
                continue
            owners = sorted(set(claimants))
            if len(owners) == 1:
                conflicts.append(f"'{name}' is declared more than once by column '{owners[0]}'")
            else:
                listed = ", ".join(f"'{owner}'" for owner in owners)
                conflicts.append(f"'{name}' is claimed by more than one column: {listed}")

        if conflicts:
            raise ValueError("Column names must have exactly one owner. " + "; ".join(sorted(conflicts)))
        return self

    def __repr__(self) -> str:
        """Return string representation of the schema."""
        cols = ", ".join(self.columns.keys())
        return f"<SchemaModel coerce={self.coerce}, on_failure={self.on_failure!r}, columns=[{cols}]>"

    # ------------------------------------------------------------------
    # Schema queries
    #
    # Purpose-specific views onto the column definitions, so consumers ask a
    # question once instead of each re-implementing the traversal. Deliberately
    # not one generic field-filter API: named queries keep call sites readable
    # and stay type-safe.
    #
    # The model is frozen, so each view is computed once per schema and reused.
    # Views returning a mapping hand back a read-only proxy, so a caller cannot
    # corrupt the cached value.
    # ------------------------------------------------------------------

    @model_validator(mode="after")
    def build_indexes(self) -> "SchemaModel":
        """Compute every derived view now, so a constructed schema is complete.

        Runs after ``verify_name_ownership``, which the indexes rely on: they
        assume every accepted name has exactly one owner.

        The views are ``cached_property``, so touching each one here populates
        it. Doing that at construction rather than on first use means the cost is
        paid once per schema instead of on some later validation run, the frozen
        object never mutates after it exists, and there is no first-call
        latency for a caller to trip over.
        """
        self._build_derived_views()
        return self

    def _build_derived_views(self) -> None:
        """Populate every derived view. Each is a ``cached_property``."""
        for view in _DERIVED_VIEWS:
            getattr(self, view)

    def __getstate__(self) -> dict[str, Any]:
        """Drop the derived views before pickling.

        They hold ``MappingProxyType`` values, which cannot be pickled, and they
        are pure functions of the schema's fields, so rebuilding is both cheap
        and guaranteed to agree.
        """
        state = super().__getstate__()
        state["__dict__"] = {k: v for k, v in state["__dict__"].items() if k not in _DERIVED_VIEWS}
        return state

    def __setstate__(self, state: dict[str, Any]) -> None:
        """Restore the schema and rebuild its derived views."""
        super().__setstate__(state)
        self._build_derived_views()

    def __deepcopy__(self, memo: dict[int, Any] | None = None) -> "SchemaModel":
        """Deep-copy the schema's fields and rebuild its derived views.

        Pydantic's default copies ``__dict__`` wholesale, which fails on the
        read-only mappings the derived views hold. Copying only the fields and
        revalidating regenerates them.
        """
        fields = {name: copy.deepcopy(getattr(self, name), memo) for name in type(self).model_fields}
        return type(self).model_validate(fields)

    @cached_property
    def resolved_columns(self) -> Mapping[str, ResolvedColumn]:
        """Every column with its inherited settings already resolved.

        Computed once, when first asked, since the schema is frozen. This is the
        shape consumers should read; ``columns`` is the authoring shape.
        """
        return MappingProxyType(
            {
                name: ResolvedColumn(
                    name=name,
                    dtype=column.dtype,
                    synonyms=tuple(column.synonyms),
                    parsers=tuple(column.parsers),
                    checks=tuple(column.checks),
                    required=column.required,
                    nullable=column.nullable,
                    coerce=self.coerce if column.coerce is None else column.coerce,
                    on_failure=self._resolve_on_failure(name),
                )
                for name, column in self.columns.items()
            }
        )

    @cached_property
    def canonical_by_cleaned_name(self) -> Mapping[str, str]:
        """Cleaned accepted name to canonical name, derived from the exact index.

        Built from `canonical_by_accepted_name` rather than as a second matcher, so
        there is one definition of what names a schema accepts.

        Raises:
            ValueError: If two accepted names clean to the same string for different
                columns, which would make a cleaned match a coin toss.
        """
        cleaned: dict[str, str] = {}
        collisions: dict[str, set[str]] = {}
        for accepted, canonical in self.canonical_by_accepted_name.items():
            key = _clean_name(accepted)
            if key in cleaned and cleaned[key] != canonical:
                collisions.setdefault(key, {cleaned[key]}).add(canonical)
            cleaned.setdefault(key, canonical)
        if collisions and self.column_matching == "cleaned":
            detail = "; ".join(f"'{key}' claimed by {sorted(names)}" for key, names in sorted(collisions.items()))
            raise ValueError(f"column_matching='cleaned' makes these names ambiguous: {detail}")
        return MappingProxyType(cleaned)

    def resolve_columns(self, physical_names: Iterable[str]) -> ColumnResolution:
        """Match a frame's physical column names against this schema.

        Pure set work over the schema's cached name index, so nothing about the
        schema is recomputed per call: intersect the physical names with the
        accepted names, reverse-map each match to its canonical column, group to
        find columns claimed by more than one name, and difference against the
        required names to find what is missing.

        Column order does not affect the result.

        Args:
            physical_names: Column names present in the data.

        Returns:
            The resolution, including any ambiguity or missing required columns.
            Inspect ``is_valid`` rather than assuming success.
        """
        index = self.canonical_by_accepted_name
        names = list(physical_names)
        matched = set(names) & self.accepted_names

        claimed: dict[str, list[str]] = {}
        for physical in sorted(matched):
            claimed.setdefault(index[physical], []).append(physical)

        if self.column_matching == "cleaned":
            # A second pass, so exact always wins. Ambiguity is judged at the winning
            # level: a column already claimed exactly ignores its cleaned candidates,
            # while a column claimed only by two cleaned names is a coin toss and
            # stays unresolved.
            cleaned_index = self.canonical_by_cleaned_name
            candidates: dict[str, list[str]] = {}
            for physical in sorted(set(names) - matched):
                canonical = cleaned_index.get(_clean_name(physical))
                if canonical is not None and canonical not in claimed:
                    candidates.setdefault(canonical, []).append(physical)
            claimed.update(candidates)

        ambiguous = {canonical: tuple(physicals) for canonical, physicals in claimed.items() if len(physicals) > 1}
        rename = {
            physicals[0]: canonical
            for canonical, physicals in claimed.items()
            if canonical not in ambiguous and physicals[0] != canonical
        }
        # Schema order, so callers and error messages stay deterministic.
        resolved = tuple(name for name in self.columns if name in claimed)
        missing_required = tuple(name for name in self.required_columns if name not in claimed)

        return ColumnResolution(
            rename=MappingProxyType(rename),
            resolved=resolved,
            missing_required=missing_required,
            ambiguous=MappingProxyType(ambiguous),
        )

    def column(self, col_name: str) -> ResolvedColumn:
        """Get one column with its inherited settings resolved.

        Args:
            col_name: Canonical column name.

        Returns:
            The resolved column.

        Raises:
            KeyError: If no column with that canonical name exists.
        """
        return self.resolved_columns[col_name]

    def _resolve_on_failure(self, col_name: str) -> OnFailureBehavior:
        """Resolve on_failure for one column, before ``resolved_columns`` exists."""
        col_schema = self.columns[col_name]
        return col_schema.on_failure if col_schema.on_failure is not None else self.on_failure

    @cached_property
    def required_columns(self) -> tuple[str, ...]:
        """Canonical names of columns that must be present in the input."""
        return tuple(name for name, column in self.resolved_columns.items() if column.required)

    @cached_property
    def non_nullable_columns(self) -> tuple[str, ...]:
        """Canonical names of columns that reject null values."""
        return tuple(name for name, column in self.resolved_columns.items() if not column.nullable)

    @cached_property
    def columns_with_parsers(self) -> tuple[str, ...]:
        """Canonical names of columns that declare at least one parser."""
        return tuple(name for name, column in self.resolved_columns.items() if column.has_parsers)

    @cached_property
    def columns_with_checks(self) -> tuple[str, ...]:
        """Canonical names of columns that declare at least one check.

        Declared checks only, matching ``columns_with_parsers``. The not-null
        constraint is generated rather than declared, so it is excluded here. Use
        ``columns_needing_check_phase`` for the wider set the phase acts on.
        """
        return tuple(name for name, column in self.resolved_columns.items() if column.has_checks)

    @cached_property
    def columns_needing_check_phase(self) -> tuple[str, ...]:
        """Canonical names of columns the check phase produces a mask for.

        A non-nullable column needs the phase even with no declared checks, since
        the not-null mask is built there.
        """
        return tuple(name for name, column in self.resolved_columns.items() if column.needs_check_phase)

    @cached_property
    def columns_to_coerce(self) -> tuple[str, ...]:
        """Canonical names of columns whose resolved coerce setting is True."""
        return tuple(name for name, column in self.resolved_columns.items() if column.coerce)

    @cached_property
    def accepted_names(self) -> frozenset[str]:
        """Every exact column name this schema accepts, canonical and synonym alike."""
        return frozenset(self.canonical_by_accepted_name)

    @cached_property
    def canonical_by_accepted_name(self) -> Mapping[str, str]:
        """Reverse index from each accepted name to the canonical column owning it.

        Name ownership is validated at construction, so every accepted name maps
        to exactly one canonical column and this index is unambiguous.

        Returns:
            Mapping of accepted name to canonical name. Canonical names map to
            themselves.
        """
        index: dict[str, str] = {}
        for canonical, column in self.columns.items():
            index[canonical] = canonical
            for synonym in column.synonyms:
                index[synonym] = canonical
        return MappingProxyType(index)

    def resolve_coerce(self, col_name: str) -> bool:
        """Resolve effective coerce setting for a column.

        Resolution order:
        1. Column coerce if set explicitly.
        2. Schema coerce as default.

        Args:
            col_name: Name of the column.

        Returns:
            Whether to coerce this column.
        """
        return self.column(col_name).coerce

    def resolve_on_failure(self, col_name: str) -> OnFailureBehavior:
        """Resolve effective on_failure for a column.

        Resolution order:
        1. Column on_failure if set explicitly.
        2. Schema on_failure as default.

        Args:
            col_name: Name of the column.

        Returns:
            Resolved on_failure behavior.
        """
        return self.column(col_name).on_failure

    @classmethod
    def from_dict(cls, data: Mapping[str, Any]) -> "SchemaModel":
        """Load a schema from a dictionary.

        Args:
            data: Dictionary representation of a schema.

        Returns:
            SchemaModel: Parsed schema model.

        Raises:
            ValueError: If validation fails.
        """
        try:
            return cls.model_validate(data)
        except PydanticValidationError as err:
            raise ValueError(f"Invalid schema configuration: {err}") from err

    @classmethod
    def from_python(cls, schema: "SchemaModel | Mapping[str, Any]") -> "SchemaModel":
        """Accept an existing SchemaModel or a dictionary defining one.

        Args:
            schema: Schema model instance or dictionary.

        Returns:
            SchemaModel: Parsed or passed-through schema.
        """
        if isinstance(schema, Mapping):
            return cls.from_dict(schema)
        return schema

    @classmethod
    def from_json(cls, content: str) -> "SchemaModel":
        """Load a schema from a JSON string.

        Args:
            content: JSON text.

        Returns:
            SchemaModel: Parsed schema model.

        Raises:
            ValueError: If JSON is invalid or schema validation fails.
        """
        try:
            data = json.loads(content)
        except json.JSONDecodeError as err:
            raise ValueError(f"Invalid JSON: {err}") from err
        return cls.from_dict(data)

    @classmethod
    def from_json_file(cls, path: str | Path) -> "SchemaModel":
        """Load a schema from a JSON file.

        Args:
            path: Path to JSON schema file.

        Returns:
            SchemaModel: Parsed schema model.

        Raises:
            ValueError: If file cannot be read or schema is invalid.
        """
        try:
            text = Path(path).read_text()
        except OSError as err:
            raise ValueError(f"Cannot read file {path}: {err}") from err
        return cls.from_json(text)

    @classmethod
    def from_yaml(cls, content: str) -> "SchemaModel":
        """Load a schema from a YAML string.

        Args:
            content: YAML text.

        Returns:
            SchemaModel: Parsed schema model.

        Raises:
            ValueError: If YAML is invalid or schema validation fails.
        """
        try:
            data = yaml.safe_load(content)
        except yaml.YAMLError as err:
            raise ValueError(f"Invalid YAML: {err}") from err
        return cls.from_dict(data)

    @classmethod
    def from_yaml_file(cls, path: str | Path) -> "SchemaModel":
        """Load a schema from a YAML file.

        Args:
            path: Path to YAML schema file.

        Returns:
            SchemaModel: Parsed schema model.

        Raises:
            ValueError: If file cannot be read or schema is invalid.
        """
        try:
            text = Path(path).read_text()
        except OSError as err:
            raise ValueError(f"Cannot read file {path}: {err}") from err
        return cls.from_yaml(text)

    @classmethod
    def from_file(cls, path: str | Path) -> "SchemaModel":
        """Load a schema from a file, auto-detecting format from extension.

        Args:
            path: Path to schema file (.json, .yaml, or .yml).

        Returns:
            SchemaModel: Parsed schema model.

        Raises:
            ValueError: If file extension is not recognized or schema is invalid.
        """
        path_obj = Path(path)
        suffix = path_obj.suffix.lower()
        if suffix == ".json":
            return cls.from_json_file(path_obj)
        if suffix in {".yaml", ".yml"}:
            return cls.from_yaml_file(path_obj)
        raise ValueError(f"Unsupported file extension '{suffix}'. Use .json, .yaml, or .yml")

    @property
    def resolved_lazy(self) -> bool:
        """Whether this run stays lazy."""
        return Config.lazy()

    @property
    def resolved_streaming_row_threshold(self) -> int:
        """The streaming threshold."""
        return Config.streaming_row_threshold()

    def verify(self, registry: Registry) -> None:
        """Check the schema against a registry, without touching any data.

        Every problem is reported at once, not just the first one found.

        Args:
            registry: Registry the schema's checks and parsers must resolve in.

        Raises:
            ConfigurationError: If any name is unresolved, any argument does not bind,
                or a column names the same check twice.
        """
        problems: list[str] = []
        pairs = (
            ("parser", registry.column_parsers, [(c, p) for c, col in self.columns.items() for p in col.parsers]),
            ("check", registry.column_checks, [(c, k) for c, col in self.columns.items() for k in col.checks]),
            ("frame parser", registry.frame_parsers, [(None, p) for p in self.frame_parsers]),
            ("frame check", registry.frame_checks, [(None, k) for k in self.frame_checks]),
        )
        for kind, sub_registry, specs in pairs:
            for column, spec in specs:
                where = f" on column '{column}'" if column else ""
                try:
                    validator = sub_registry.get(spec.name)
                except KeyError:
                    available = ", ".join(sorted(sub_registry.list_names())) or "none"
                    problems.append(f"{kind} '{spec.name}'{where} is not registered. Available: {available}")
                    continue
                try:
                    validator.validate_args(**spec.args)
                except (TypeError, ValueError) as e:
                    problems.append(f"{kind} '{spec.name}'{where} has invalid arguments: {e}")

        for column, col_schema in self.columns.items():
            seen: set[str] = set()
            for check in col_schema.checks:
                if check.name in seen:
                    problems.append(
                        f"check '{check.name}' is declared more than once on column '{column}'. "
                        "Each check name may appear once per column, since the error report is "
                        "keyed on (column, check)."
                    )
                seen.add(check.name)

        if problems:
            raise ConfigurationError("Schema does not verify against the registry:\n  " + "\n  ".join(problems))

    def validate(  # ty: ignore[invalid-method-override]
        self,
        df: pl.DataFrame | pl.LazyFrame,
        registry: Registry,
        *,
        error_report_config: "ErrorReportConfig | None" = None,
        lazy: bool | None = None,
        pipeline: "ValidationPipeline | None" = None,
    ) -> "ValidationResult":
        """Validate a DataFrame against this schema.

        Args:
            df: DataFrame to validate.
            registry: Validator registry with parsers and checks.
            error_report_config: How much detail `result.errors` carries.
            lazy: Return a LazyFrame (True) or a DataFrame (False). If None, reads
                `nyctea.Config`.
            pipeline: Phases to run. If None, the phases the schema needs are built
                for you. Pass one to reorder phases or add your own.

        Returns:
            ValidationResult with validated data, errors, and report.

        Raises:
            ConfigurationError: If the schema does not verify against the registry. This
                runs before any data is read, so it precedes both errors below.
            nyctea.ValidationError: If the input does not match the schema's structure,
                because a required column is missing from it or a name resolves
                ambiguously. Named in full because this module binds pydantic's
                `ValidationError` under an alias.
            nyctea.PipelineError: If a check, parser, coercion or nullability failure is
                set to `on_failure="raise"`, or if a phase fails for any other reason.

        Example:
            >>> from nyctea.validators.registry import Registry
            >>> schema = SchemaModel.from_yaml("schema.yaml")
            >>> registry = Registry()
            >>> # ... register validators ...
            >>> result = schema.validate(df, registry)
            >>> print(result.report.summary())
        """
        from nyctea.engine.validator import DataValidator

        validator = DataValidator(self, registry, pipeline=pipeline)
        return validator.validate(df, error_report_config=error_report_config, lazy=lazy)