patito 0.7.0__py3-none-any.whl → 0.8.2__py3-none-any.whl

patito/pydantic.py

@@ -3,58 +3,46 @@
 from __future__ import annotations
 
 import itertools
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping, Sequence
 from datetime import date, datetime, time, timedelta
-from functools import partial
 from inspect import getfullargspec
 from typing import (
     TYPE_CHECKING,
     Any,
     ClassVar,
-    Dict,
-    FrozenSet,
-    Generic,
-    List,
     Literal,
-    Mapping,
     Optional,
-    Sequence,
-    Tuple,
-    Type,
     TypeVar,
-    Union,
     cast,
     get_args,
 )
+from zoneinfo import ZoneInfo
 
 import polars as pl
 from polars.datatypes import DataType, DataTypeClass
 from pydantic import (  # noqa: F401
     BaseModel,
     create_model,
-    field_serializer,
     fields,
 )
 from pydantic._internal._model_construction import (
     ModelMetaclass as PydanticModelMetaclass,
 )
-from zoneinfo import ZoneInfo
 
-from patito._pydantic.column_info import CI, ColumnInfo
+from patito._pydantic.column_info import ColumnInfo
 from patito._pydantic.dtypes import (
     default_dtypes_for_model,
-    dtype_from_string,
     is_optional,
     valid_dtypes_for_model,
     validate_annotation,
     validate_polars_dtype,
 )
 from patito._pydantic.schema import column_infos_for_model, schema_for_model
-from patito.polars import DataFrame, LazyFrame
+from patito.polars import DataFrame, LazyFrame, ModelGenerator
 from patito.validators import validate
 
 try:
-    import pandas as pd
+    import pandas as pd  # type: ignore
 
     _PANDAS_AVAILABLE = True
 except ImportError:
@@ -68,16 +56,14 @@ if TYPE_CHECKING:
 ModelType = TypeVar("ModelType", bound="Model")
 
 
-class ModelMetaclass(PydanticModelMetaclass, Generic[CI]):
+class ModelMetaclass(PydanticModelMetaclass):
     """Metaclass used by patito.Model.
 
     Responsible for setting any relevant model-dependent class properties.
     """
 
-    column_info_class: ClassVar[Type[ColumnInfo]] = ColumnInfo
-
     if TYPE_CHECKING:
-        model_fields: ClassVar[Dict[str, fields.FieldInfo]]
+        model_fields: ClassVar[dict[str, fields.FieldInfo]]
 
     def __init__(cls, name: str, bases: tuple, clsdict: dict, **kwargs) -> None:
         """Construct new patito model.
@@ -90,27 +76,31 @@ class ModelMetaclass(PydanticModelMetaclass, Generic[CI]):
 
         """
         super().__init__(name, bases, clsdict, **kwargs)
-        # Add a custom subclass of patito.DataFrame to the model class,
-        # where .set_model() has been implicitly set.
-        cls.DataFrame = DataFrame._construct_dataframe_model_class(
-            model=cls,  # type: ignore
+        NewDataFrame = type(
+            f"{cls.__name__}DataFrame",
+            (DataFrame,),
+            {"model": cls},
         )
-        # Similarly for LazyFrame
-        cls.LazyFrame = LazyFrame._construct_lazyframe_model_class(
-            model=cls,  # type: ignore
+        cls.DataFrame: type[DataFrame[cls]] = NewDataFrame  # type: ignore
+
+        NewLazyFrame = type(
+            f"{cls.__name__}LazyFrame",
+            (LazyFrame,),
+            {"model": cls},
         )
+        cls.LazyFrame: type[LazyFrame[cls]] = NewLazyFrame  # type: ignore
 
     def __hash__(self) -> int:
         """Return hash of the model class."""
         return super().__hash__()
 
     @property
-    def column_infos(cls: Type[ModelType]) -> Mapping[str, ColumnInfo]:
+    def column_infos(cls: type[Model]) -> Mapping[str, ColumnInfo]:
         """Return column information for the model."""
         return column_infos_for_model(cls)
 
     @property
-    def model_schema(cls: Type[ModelType]) -> Mapping[str, Mapping[str, Any]]:
+    def model_schema(cls: type[Model]) -> Mapping[str, Mapping[str, Any]]:
         """Return schema properties where definition references have been resolved.
 
         Returns:
@@ -126,7 +116,7 @@ class ModelMetaclass(PydanticModelMetaclass, Generic[CI]):
         return schema_for_model(cls)
 
     @property
-    def columns(cls: Type[ModelType]) -> List[str]:
+    def columns(cls: type[Model]) -> list[str]:
         """Return the name of the dataframe columns specified by the fields of the model.
 
         Returns:
@@ -145,7 +135,7 @@ class ModelMetaclass(PydanticModelMetaclass, Generic[CI]):
         return list(cls.model_fields.keys())
 
     @property
-    def dtypes(cls: Type[ModelType]) -> dict[str, DataTypeClass | DataType]:
+    def dtypes(cls: type[Model]) -> dict[str, DataTypeClass | DataType]:
         """Return the polars dtypes of the dataframe.
 
         Unless Field(dtype=...) is specified, the highest signed column dtype
@@ -169,8 +159,8 @@ class ModelMetaclass(PydanticModelMetaclass, Generic[CI]):
 
     @property
     def valid_dtypes(
-        cls: Type[ModelType],
-    ) -> Mapping[str, FrozenSet[DataTypeClass | DataType]]:
+        cls: type[Model],
+    ) -> Mapping[str, frozenset[DataTypeClass | DataType]]:
         """Return a list of polars dtypes which Patito considers valid for each field.
 
         The first item of each list is the default dtype chosen by Patito.
@@ -186,7 +176,7 @@ class ModelMetaclass(PydanticModelMetaclass, Generic[CI]):
         return valid_dtypes_for_model(cls)
 
     @property
-    def defaults(cls: Type[ModelType]) -> dict[str, Any]:
+    def defaults(cls: type[Model]) -> dict[str, Any]:
         """Return default field values specified on the model.
 
         Returns:
@@ -211,7 +201,7 @@ class ModelMetaclass(PydanticModelMetaclass, Generic[CI]):
         }
 
     @property
-    def non_nullable_columns(cls: Type[ModelType]) -> set[str]:
+    def non_nullable_columns(cls: type[Model]) -> set[str]:
         """Return names of those columns that are non-nullable in the schema.
 
         Returns:
@@ -235,12 +225,12 @@ class ModelMetaclass(PydanticModelMetaclass, Generic[CI]):
             for k in cls.columns
             if not (
                 is_optional(cls.model_fields[k].annotation)
-                or cls.model_fields[k].annotation == type(None)
+                or cls.model_fields[k].annotation is type(None)
             )
         )
 
     @property
-    def nullable_columns(cls: Type[ModelType]) -> set[str]:
+    def nullable_columns(cls: type[Model]) -> set[str]:
         """Return names of those columns that are nullable in the schema.
 
         Returns:
@@ -262,7 +252,7 @@ class ModelMetaclass(PydanticModelMetaclass, Generic[CI]):
         return set(cls.columns) - cls.non_nullable_columns
 
     @property
-    def unique_columns(cls: Type[ModelType]) -> set[str]:
+    def unique_columns(cls: type[Model]) -> set[str]:
         """Return columns with uniqueness constraint.
 
         Returns:
@@ -285,7 +275,7 @@ class ModelMetaclass(PydanticModelMetaclass, Generic[CI]):
         return {column for column in cls.columns if infos[column].unique}
 
     @property
-    def derived_columns(cls: Type[ModelType]) -> set[str]:
+    def derived_columns(cls: type[Model]) -> set[str]:
         """Return set of columns which are derived from other columns."""
         infos = cls.column_infos
         return {
@@ -297,7 +287,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
     """Custom pydantic class for representing table schema and constructing rows."""
 
     @classmethod
-    def validate_schema(cls: Type[ModelType]):
+    def validate_schema(cls: type[ModelType]):
         """Users should run this after defining or edit a model. We withhold the checks at model definition time to avoid expensive queries of the model schema."""
         for column in cls.columns:
             col_info = cls.column_infos[column]
@@ -311,8 +301,8 @@ class Model(BaseModel, metaclass=ModelMetaclass):
 
     @classmethod
     def from_row(
-        cls: Type[ModelType],
-        row: Union["pd.DataFrame", pl.DataFrame],
+        cls: type[ModelType],
+        row: pd.DataFrame | pl.DataFrame,
         validate: bool = True,
     ) -> ModelType:
         """Represent a single data frame row as a Patito model.
@@ -340,6 +330,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
             >>> df = pl.DataFrame(
             ...     [["1", "product name", "1.22"]],
             ...     schema=["product_id", "name", "price"],
+            ...     orient="row",
             ... )
             >>> Product.from_row(df)
             Product(product_id=1, name='product name', price=1.22)
@@ -359,7 +350,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
 
     @classmethod
     def _from_polars(
-        cls: Type[ModelType],
+        cls: type[ModelType],
         dataframe: pl.DataFrame,
         validate: bool = True,
     ) -> ModelType:
@@ -392,6 +383,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
             >>> df = pl.DataFrame(
             ...     [["1", "product name", "1.22"]],
             ...     schema=["product_id", "name", "price"],
+            ...     orient="row",
             ... )
             >>> Product._from_polars(df)
             Product(product_id=1, name='product name', price=1.22)
@@ -419,13 +411,13 @@ class Model(BaseModel, metaclass=ModelMetaclass):
 
     @classmethod
     def validate(
-        cls,
-        dataframe: Union["pd.DataFrame", pl.DataFrame],
-        columns: Optional[Sequence[str]] = None,
+        cls: type[ModelType],
+        dataframe: pd.DataFrame | pl.DataFrame,
+        columns: Sequence[str] | None = None,
         allow_missing_columns: bool = False,
         allow_superfluous_columns: bool = False,
-        **kwargs,
-    ) -> None:
+        drop_superfluous_columns: bool = False,
+    ) -> DataFrame[ModelType]:
         """Validate the schema and content of the given dataframe.
 
         Args:
@@ -434,10 +426,11 @@ class Model(BaseModel, metaclass=ModelMetaclass):
                 of the dataframe will be validated.
             allow_missing_columns: If True, missing columns will not be considered an error.
             allow_superfluous_columns: If True, additional columns will not be considered an error.
-            **kwargs: Additional keyword arguments to be passed to the validation
+            drop_superfluous_columns: If True, columns not present in the model will be
+                dropped from the resulting dataframe.
 
         Returns:
-            ``None``:
+            DataFrame: A patito DataFrame containing the validated data.
 
         Raises:
             patito.exceptions.DataFrameValidationError: If the given dataframe does not match
@@ -479,15 +472,35 @@ class Model(BaseModel, metaclass=ModelMetaclass):
             columns=columns,
             allow_missing_columns=allow_missing_columns,
             allow_superfluous_columns=allow_superfluous_columns,
-            **kwargs,
+            drop_superfluous_columns=drop_superfluous_columns,
         )
+        return cls.DataFrame(dataframe)
+
+    @classmethod
+    def iter_models(
+        cls: type[ModelType], dataframe: pd.DataFrame | pl.DataFrame
+    ) -> ModelGenerator[ModelType]:
+        """Validate the dataframe and iterate over the rows, yielding Patito models.
+
+        Args:
+            dataframe: Polars or pandas DataFrame to be validated.
+
+        Returns:
+            ListableIterator: An iterator of patito models over the validated data.
+
+        Raises:
+            patito.exceptions.DataFrameValidationError: If the given dataframe does not match
+                the given schema.
+
+        """
+        return cls.DataFrame(dataframe).iter_models()
 
     @classmethod
     def example_value(  # noqa: C901
         cls,
-        field: Optional[str] = None,
-        properties: Optional[Dict[str, Any]] = None,
-    ) -> Union[date, datetime, time, timedelta, float, int, str, None, Mapping, List]:
+        field: str | None = None,
+        properties: dict[str, Any] | None = None,
+    ) -> date | datetime | time | timedelta | float | int | str | None | Mapping | list:
         """Return a valid example value for the given model field.
 
         Args:
@@ -529,7 +542,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
             properties = cls._schema_properties()[field]
             info = cls.column_infos[field]
         else:
-            info = cls.column_info_class()
+            info = ColumnInfo()
         properties = properties or {}
 
         if "type" in properties:
@@ -611,10 +624,10 @@ class Model(BaseModel, metaclass=ModelMetaclass):
                 return date(year=1970, month=1, day=1)
             elif "format" in properties and properties["format"] == "date-time":
                 if "column_info" in properties:
-                    dtype_str = properties["column_info"]["dtype"]
-                    dtype = dtype_from_string(dtype_str)
+                    ci = ColumnInfo.model_validate_json(properties["column_info"])
+                    dtype = ci.dtype
                     if getattr(dtype, "time_zone", None) is not None:
-                        tzinfo = ZoneInfo(dtype.time_zone)
+                        tzinfo = ZoneInfo(dtype.time_zone)  # type: ignore
                     else:
                         tzinfo = None
                     return datetime(year=1970, month=1, day=1, tzinfo=tzinfo)
@@ -650,7 +663,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
 
     @classmethod
     def example(
-        cls: Type[ModelType],
+        cls: type[ModelType],
         **kwargs: Any,  # noqa: ANN401
     ) -> ModelType:
         """Produce model instance with filled dummy data for all unspecified fields.
@@ -702,10 +715,10 @@ class Model(BaseModel, metaclass=ModelMetaclass):
 
     @classmethod
     def pandas_examples(
-        cls: Type[ModelType],
-        data: Union[dict, Iterable],
-        columns: Optional[Iterable[str]] = None,
-    ) -> "pd.DataFrame":
+        cls: type[ModelType],
+        data: dict | Iterable,
+        columns: Iterable[str] | None = None,
+    ) -> pd.DataFrame:
         """Generate dataframe with dummy data for all unspecified columns.
 
         Offers the same API as the pandas.DataFrame constructor.
@@ -772,10 +785,10 @@ class Model(BaseModel, metaclass=ModelMetaclass):
 
     @classmethod
     def examples(
-        cls: Type[ModelType],
-        data: Optional[Union[dict, Iterable]] = None,
-        columns: Optional[Iterable[str]] = None,
-    ) -> "patito.polars.DataFrame":
+        cls: type[ModelType],
+        data: dict | Iterable | None = None,
+        columns: Iterable[str] | None = None,
+    ) -> patito.polars.DataFrame:
         """Generate polars dataframe with dummy data for all unspecified columns.
 
         This constructor accepts the same data format as polars.DataFrame.
@@ -844,7 +857,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
         if wrong_columns:
             raise TypeError(f"{cls.__name__} does not contain fields {wrong_columns}!")
 
-        series: List[Union[pl.Series, pl.Expr]] = []
+        series: list[pl.Series | pl.Expr] = []
         unique_series = []
         for column_name, dtype in cls.dtypes.items():
             if column_name not in kwargs:
@@ -872,10 +885,10 @@ class Model(BaseModel, metaclass=ModelMetaclass):
 
     @classmethod
     def join(
-        cls: Type["Model"],
-        other: Type["Model"],
+        cls: type[Model],
+        other: type[Model],
         how: Literal["inner", "left", "outer", "asof", "cross", "semi", "anti"],
-    ) -> Type["Model"]:
+    ) -> type[Model]:
         """Dynamically create a new model compatible with an SQL Join operation.
 
         For instance, ``ModelA.join(ModelB, how="left")`` will create a model containing
@@ -920,7 +933,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
         if how in {"semi", "anti"}:
             return cls
 
-        kwargs: Dict[str, Any] = {}
+        kwargs: dict[str, Any] = {}
         for model, nullable_methods in (
             (cls, {"outer"}),
             (other, {"left", "outer", "asof"}),
@@ -940,9 +953,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
         )
 
     @classmethod
-    def select(
-        cls: Type[ModelType], fields: Union[str, Iterable[str]]
-    ) -> Type["Model"]:
+    def select(cls: type[ModelType], fields: str | Iterable[str]) -> type[Model]:
         """Create a new model consisting of only a subset of the model fields.
 
         Args:
@@ -984,7 +995,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
         )
 
     @classmethod
-    def drop(cls: Type[ModelType], name: Union[str, Iterable[str]]) -> Type["Model"]:
+    def drop(cls: type[ModelType], name: str | Iterable[str]) -> type[Model]:
         """Return a new model where one or more fields are excluded.
 
         Args:
@@ -1023,7 +1034,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
         )
 
     @classmethod
-    def prefix(cls: Type[ModelType], prefix: str) -> Type["Model"]:
+    def prefix(cls: type[ModelType], prefix: str) -> type[Model]:
         """Return a new model where all field names have been prefixed.
 
         Args:
@@ -1049,7 +1060,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
         )
 
     @classmethod
-    def suffix(cls: Type[ModelType], suffix: str) -> Type["Model"]:
+    def suffix(cls: type[ModelType], suffix: str) -> type[Model]:
         """Return a new model where all field names have been suffixed.
 
         Args:
@@ -1076,7 +1087,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
         )
 
     @classmethod
-    def rename(cls: Type[ModelType], mapping: Dict[str, str]) -> Type["Model"]:
+    def rename(cls: type[ModelType], mapping: dict[str, str]) -> type[Model]:
         """Return a new model class where the specified fields have been renamed.
 
         Args:
@@ -1117,9 +1128,9 @@ class Model(BaseModel, metaclass=ModelMetaclass):
 
     @classmethod
     def with_fields(
-        cls: Type[ModelType],
+        cls: type[ModelType],
         **field_definitions: Any,  # noqa: ANN401
-    ) -> Type["Model"]:
+    ) -> type[Model]:
         """Return a new model class where the given fields have been added.
 
         Args:
@@ -1152,11 +1163,11 @@ class Model(BaseModel, metaclass=ModelMetaclass):
         )
 
     @classmethod
-    def _schema_properties(cls: Type[ModelType]) -> Mapping[str, Any]:
+    def _schema_properties(cls: type[ModelType]) -> Mapping[str, Any]:
         return cls.model_schema["properties"]
 
     @classmethod
-    def _update_dfn(cls, annotation: Any, schema: Dict[str, Any]) -> None:
+    def _update_dfn(cls, annotation: Any, schema: dict[str, Any]) -> None:
         try:
             if issubclass(annotation, Model) and annotation.__name__ != cls.__name__:
                 schema["$defs"][annotation.__name__] = annotation.model_schema
@@ -1165,10 +1176,10 @@ class Model(BaseModel, metaclass=ModelMetaclass):
 
     @classmethod
     def _derive_model(
-        cls: Type[ModelType],
+        cls: type[ModelType],
         model_name: str,
-        field_mapping: Dict[str, Any],
-    ) -> Type["Model"]:
+        field_mapping: dict[str, Any],
+    ) -> type[Model]:
         """Derive a new model with new field definitions.
 
         Args:
@@ -1209,7 +1220,7 @@ class Model(BaseModel, metaclass=ModelMetaclass):
     def _derive_field(
         field: fields.FieldInfo,
         make_nullable: bool = False,
-    ) -> Tuple[Type | None, fields.FieldInfo]:
+    ) -> tuple[type | None, fields.FieldInfo]:
         field_type = field.annotation
         default = field.default
         extra_attrs = {
@@ -1241,8 +1252,8 @@ FIELD_KWARGS = getfullargspec(fields.Field)
 # Helper function for patito Field.
 
 
-def FieldCI(
-    column_info: Type[ColumnInfo], *args: Any, **kwargs: Any
+def Field(
+    *args: Any, **kwargs: Any
 ) -> Any:  # annotate with Any to make the downstream type annotations happy
     """Annotate model field with additional type and validation information.
 
@@ -1254,6 +1265,7 @@ def FieldCI(
     can be read with the below examples.
 
     Args:
+        allow_missing (bool): Column may be missing.
         column_info: (Type[ColumnInfo]): ColumnInfo object to pass args to.
         constraints (Union[polars.Expression, List[polars.Expression]): A single
             constraint or list of constraints, expressed as a polars expression objects.
@@ -1313,7 +1325,7 @@ def FieldCI(
             Polars dtype Int64 does not match model field type. (type=type_error.columndtype)
 
     """
-    ci = column_info(**kwargs)
+    ci = ColumnInfo(**kwargs)
     for field in ci.model_fields_set:
         kwargs.pop(field)
     if kwargs.pop("modern_kwargs_only", True):
@@ -1322,11 +1334,12 @@ def FieldCI(
                 raise ValueError(
                     f"unexpected kwarg {kwarg}={kwargs[kwarg]}.  Add modern_kwargs_only=False to ignore"
                 )
+    ci_json = ci.model_dump_json()
+    existing_json_schema_extra = kwargs.pop("json_schema_extra", {})
+    merged_json_schema_extra = {**existing_json_schema_extra, "column_info": ci_json}
+
     return fields.Field(
         *args,
-        json_schema_extra={"column_info": ci},
+        json_schema_extra=merged_json_schema_extra,
         **kwargs,
     )
-
-
-Field = partial(FieldCI, column_info=ColumnInfo)