Schema Models¶

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

    model_config = ConfigDict(extra="forbid")

    lazy: bool = Field(
        True,
        description="Whether to use Polars lazy execution during validation",
    )

    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': value becomes null\n"
            "- 'ignore': coercion nulls forced by dtype, check failures 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")

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

    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.
        """
        col_schema = self.columns[col_name]
        if col_schema.coerce is not None:
            return col_schema.coerce
        return self.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.
        3. Guard: on_failure=null requires nullable=True. Non-nullable columns
           fall back to raise.

        Args:
            col_name: Name of the column.

        Returns:
            Resolved on_failure behavior.
        """
        col_schema = self.columns[col_name]

        behavior = col_schema.on_failure if col_schema.on_failure is not None else self.on_failure

        # Guard: can't nullify non-nullable columns
        if behavior == "null" and not col_schema.nullable:
            return "raise"

        return behavior

    @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 ValidationError 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:
            ImportError: If PyYAML is not installed.
            ValueError: If YAML is invalid or schema validation fails.
        """
        cls._ensure_yaml()
        assert yaml is not None  # for type checkers
        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:
            ImportError: If PyYAML is not installed.
            ValueError: If file cannot be read or schema is invalid.
        """
        cls._ensure_yaml()
        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")

    @staticmethod
    def _ensure_yaml() -> None:
        """Raise if PyYAML is unavailable."""
        if yaml is None:
            raise ImportError("PyYAML is required for YAML schema loading. Install with: pip install nyctea[yaml]")

    def validate(  # ty: ignore[invalid-method-override]
        self,
        df: pl.DataFrame | pl.LazyFrame,
        registry: "Registry",
        **kwargs: Any,
    ) -> "ValidationResult":
        """Validate a DataFrame against this schema.

        This is the primary API for validation using the new validator-based
        pipeline architecture.

        Args:
            df: DataFrame to validate.
            registry: Validator registry with parsers and checks.
            **kwargs: Additional validation options passed to SchemaValidator.

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

        Raises:
            ValidationError: If validation fails in strict mode.
            PipelineError: If pipeline execution fails.

        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.schema.validator import SchemaValidator

        validator = SchemaValidator(self, registry)
        return validator.validate(df, **kwargs)

    def create_validator(
        self,
        registry: "Registry",
        pipeline: "ValidationPipeline | None" = None,
    ) -> "SchemaValidator":
        """Create a SchemaValidator for this schema.

        This factory method allows you to create a validator and customize
        its pipeline before running validation.

        Args:
            registry: Validator registry with parsers and checks.
            pipeline: Custom pipeline (if None, creates from schema).

        Returns:
            SchemaValidator instance.

        Example:
            >>> from nyctea.validators.registry import Registry
            >>> schema = SchemaModel.from_yaml("schema.yaml")
            >>> registry = Registry()
            >>> # ... register validators ...
            >>> validator = schema.create_validator(registry)
            >>> # Customize pipeline
            >>> validator.pipeline.add_phase(MyCustomPhase(), after="column_parsing")
            >>> result = validator.validate(df)
        """
        from nyctea.schema.validator import SchemaValidator

        return SchemaValidator(self, registry, pipeline)

class ColumnSchema(BaseModel):
    """Schema for a single column."""

    model_config = ConfigDict(extra="forbid")

    dtype: str = Field(..., description="The final enforced dtype after validation")
    synonyms: list[str] = Field(default_factory=list, description="Allowed alternative names for this column")

    parsers: list[Parser] = Field(
        default_factory=list,
        description="List of parser functions applied before checking",
    )

    checks: list[Check] = Field(
        default_factory=list,
        description="List of checks applied independently on the parsed column",
    )

    required: bool = Field(
        True,
        description="Whether this column must be present in the input",
    )

    nullable: bool = Field(
        False,
        description="Whether null values are allowed in this column",
    )

    coerce: bool | None = Field(
        None,
        description="Whether to coerce this column to its dtype. None inherits from schema.",
    )

    on_failure: OnFailureBehavior | None = Field(
        None,
        description=(
            "How to handle coercion/check failures for this column:\n"
            "- 'raise': error, stop\n"
            "- 'null': value becomes null (requires nullable=True)\n"
            "- 'ignore': coercion nulls forced by dtype, check failures kept and reported\n"
            "- None: inherit from schema on_failure"
        ),
    )

    @field_validator("dtype")
    @classmethod
    def validate_dtype(cls, v: str) -> str:
        """Validate that dtype is a valid Polars dtype."""
        if not hasattr(pl, v):
            raise ValueError(f"'{v}' is not a valid Polars dtype")
        dtype_obj = getattr(pl, v)
        if not isinstance(dtype_obj, type) or not issubclass(dtype_obj, pl.DataType):
            raise TypeError(f"'{v}' is not a valid Polars DataType")
        return v

    @model_validator(mode="after")
    def validate_on_failure_nullable_consistency(self) -> "ColumnSchema":
        """Ensure on_failure='null' requires nullable=True."""
        if self.on_failure == "null" and not self.nullable:
            raise ValueError(
                "on_failure='null' requires nullable=True. Cannot nullify failures in a non-nullable column."
            )
        return self

class Parser(BaseModel):
    """Configuration for a column-level parser."""

    model_config = ConfigDict(extra="forbid")

    name: str = Field(..., description="Name of the parse function to apply")
    args: dict[str, Any] = Field(default_factory=dict, description="Arguments passed to the parse function")

class Check(BaseModel):
    """Configuration for a column-level check."""

    model_config = ConfigDict(extra="forbid")

    name: str = Field(..., description="Name of the check function to apply")
    args: dict[str, Any] = Field(default_factory=dict, description="Arguments passed to the check function")

class FrameParser(BaseModel):
    """Configuration for a frame-level parser."""

    model_config = ConfigDict(extra="forbid")

    name: str = Field(..., description="Name of the frame-level parser to apply")
    args: dict[str, Any] = Field(default_factory=dict, description="Arguments passed to the frame parser")

class FrameCheck(BaseModel):
    """Configuration for a frame-level check."""

    model_config = ConfigDict(extra="forbid")

    name: str = Field(..., description="Name of the frame-level check to apply")
    args: dict[str, Any] = Field(default_factory=dict, description="Arguments passed to the frame check")