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

patito/exceptions.py

@@ -1,15 +1,10 @@
 """Exceptions used by patito."""
 
+from collections.abc import Generator, Sequence
 from typing import (
     TYPE_CHECKING,
     Any,
-    Dict,
-    Generator,
-    List,
     Optional,
-    Sequence,
-    Tuple,
-    Type,
     TypedDict,
     Union,
 )
@@ -19,7 +14,7 @@ from patito._pydantic.repr import Representation
 if TYPE_CHECKING:
     from pydantic import BaseModel
 
-    Loc = Tuple[Union[int, str], ...]
+    Loc = tuple[Union[int, str], ...]
 
     class _ErrorDictRequired(TypedDict):
         loc: Loc
@@ -27,7 +22,7 @@ if TYPE_CHECKING:
         type: str
 
     class ErrorDict(_ErrorDictRequired, total=False):
-        ctx: Dict[str, Any]
+        ctx: dict[str, Any]
 
     from patito._pydantic.repr import ReprArgs
 
@@ -67,13 +62,13 @@ class DataFrameValidationError(Representation, ValueError):
 
     __slots__ = "raw_errors", "model", "_error_cache"
 
-    def __init__(self, errors: Sequence[ErrorList], model: Type["BaseModel"]) -> None:
+    def __init__(self, errors: Sequence[ErrorList], model: type["BaseModel"]) -> None:
         """Create a dataframe validation error."""
         self.raw_errors = errors
         self.model = model
-        self._error_cache: Optional[List["ErrorDict"]] = None
+        self._error_cache: Optional[list[ErrorDict]] = None
 
-    def errors(self) -> List["ErrorDict"]:
+    def errors(self) -> list["ErrorDict"]:
         """Get list of errors."""
         if self._error_cache is None:
             self._error_cache = list(flatten_errors(self.raw_errors))
@@ -93,7 +88,7 @@ class DataFrameValidationError(Representation, ValueError):
         return [("model", self.model.__name__), ("errors", self.errors())]
 
 
-def display_errors(errors: List["ErrorDict"]) -> str:
+def display_errors(errors: list["ErrorDict"]) -> str:
     return "\n".join(
         f'{_display_error_loc(e)}\n  {e["msg"]} ({_display_error_type_and_ctx(e)})'
         for e in errors
@@ -142,7 +137,7 @@ def error_dict(exc: Exception, loc: "Loc") -> "ErrorDict":
     else:
         msg = str(exc)
 
-    d: "ErrorDict" = {"loc": loc, "msg": msg, "type": type_}
+    d: ErrorDict = {"loc": loc, "msg": msg, "type": type_}
 
     if ctx:
         d["ctx"] = ctx
@@ -150,10 +145,10 @@ def error_dict(exc: Exception, loc: "Loc") -> "ErrorDict":
     return d
 
 
-_EXC_TYPE_CACHE: Dict[Type[Exception], str] = {}
+_EXC_TYPE_CACHE: dict[type[Exception], str] = {}
 
 
-def get_exc_type(cls: Type[Exception]) -> str:
+def get_exc_type(cls: type[Exception]) -> str:
     # slightly more efficient than using lru_cache since we don't need to worry about the cache filling up
     try:
         return _EXC_TYPE_CACHE[cls]
@@ -163,7 +158,7 @@ def get_exc_type(cls: Type[Exception]) -> str:
         return r
 
 
-def _get_exc_type(cls: Type[Exception]) -> str:
+def _get_exc_type(cls: type[Exception]) -> str:
     if issubclass(cls, AssertionError):
         return "assertion_error"
 

patito/polars.py

@@ -2,20 +2,13 @@
 
 from __future__ import annotations
 
+from collections.abc import Collection, Iterable, Iterator, Sequence
 from typing import (
     TYPE_CHECKING,
     Any,
-    Collection,
-    Dict,
     Generic,
-    Iterable,
     Literal,
-    Optional,
-    Sequence,
-    Tuple,
-    Type,
     TypeVar,
-    Union,
     cast,
 )
 
@@ -31,63 +24,110 @@ if TYPE_CHECKING:
 
     from patito.pydantic import Model
 
-
 DF = TypeVar("DF", bound="DataFrame")
 LDF = TypeVar("LDF", bound="LazyFrame")
 ModelType = TypeVar("ModelType", bound="Model")
 OtherModelType = TypeVar("OtherModelType", bound="Model")
+T = TypeVar("T")
+
+
+class ModelGenerator(Iterator[ModelType], Generic[ModelType]):
+    """An iterator that can be converted to a list."""
+
+    def __init__(self, iterator: Iterator[ModelType]) -> None:
+        """Construct a ModelGenerator from an iterator."""
+        self._iterator = iterator
+
+    def to_list(self) -> list[ModelType]:
+        """Convert iterator to list."""
+        return list(self)
+
+    def __next__(self) -> ModelType:  # noqa: D105
+        return next(self._iterator)
+
+    def __iter__(self) -> Iterator[ModelType]:  # noqa: D105
+        return self
 
 
 class LazyFrame(pl.LazyFrame, Generic[ModelType]):
     """LazyFrame class associated to DataFrame."""
 
-    model: Type[ModelType]
+    model: type[ModelType]
 
-    @classmethod
-    def _construct_lazyframe_model_class(
-        cls: Type[LDF], model: Optional[Type[ModelType]]
-    ) -> Type[LazyFrame[ModelType]]:
-        """Return custom LazyFrame sub-class where LazyFrame.model is set.
+    def set_model(self, model: type[OtherModelType]) -> LazyFrame[OtherModelType]:
+        """Associate a given patito ``Model`` with the dataframe.
 
-        Can be used to construct a LazyFrame class where
-        DataFrame.set_model(model) is implicitly invoked at collection.
+        The model schema is used by methods that depend on a model being associated with
+        the given dataframe such as :ref:`DataFrame.validate() <DataFrame.validate>`
+        and :ref:`DataFrame.get() <DataFrame.get>`.
+
+        ``DataFrame(...).set_model(Model)`` is equivalent with ``Model.DataFrame(...)``.
 
         Args:
-            model: A patito model which should be used to validate the final dataframe.
-                If None is provided, the regular LazyFrame class will be returned.
+            model (Model): Sub-class of ``patito.Model`` declaring the schema of the
+                dataframe.
 
         Returns:
-            A custom LazyFrame model class where LazyFrame.model has been correctly
-                "hard-coded" to the given model.
+            DataFrame[Model]: Returns the same dataframe, but with an attached model
+            that is required for certain model-specific dataframe methods to work.
 
-        """
-        if model is None:
-            return cls
+        Examples:
+            >>> from typing_extensions import Literal
+            >>> import patito as pt
+            >>> import polars as pl
+            >>> class SchoolClass(pt.Model):
+            ...     year: int = pt.Field(dtype=pl.UInt16)
+            ...     letter: Literal["A", "B"] = pt.Field(dtype=pl.Categorical)
+            ...
+            >>> classes = pt.DataFrame(
+            ...     {"year": [1, 1, 2, 2], "letter": list("ABAB")}
+            ... ).set_model(SchoolClass)
+            >>> classes
+            shape: (4, 2)
+            ┌──────┬────────┐
+            │ year ┆ letter │
+            │ ---  ┆ ---    │
+            │ i64  ┆ str    │
+            ╞══════╪════════╡
+            │ 1    ┆ A      │
+            │ 1    ┆ B      │
+            │ 2    ┆ A      │
+            │ 2    ┆ B      │
+            └──────┴────────┘
+            >>> casted_classes = classes.cast()
+            >>> casted_classes
+            shape: (4, 2)
+            ┌──────┬────────┐
+            │ year ┆ letter │
+            │ ---  ┆ ---    │
+            │ u16  ┆ cat    │
+            ╞══════╪════════╡
+            │ 1    ┆ A      │
+            │ 1    ┆ B      │
+            │ 2    ┆ A      │
+            │ 2    ┆ B      │
+            └──────┴────────┘
+            >>> casted_classes.validate()
 
-        new_class = type(
-            f"{model.__name__}LazyFrame",
-            (cls,),
-            {"model": model},
-        )
-        return new_class
+        """
+        return model.LazyFrame._from_pyldf(self._ldf)  # type: ignore
 
     def collect(
         self,
         *args,
         **kwargs,
-    ) -> "DataFrame[ModelType]":  # noqa: DAR101, DAR201
+    ) -> DataFrame[ModelType]:  # noqa: DAR101, DAR201
         """Collect into a DataFrame.
 
         See documentation of polars.DataFrame.collect for full description of
         parameters.
         """
         background = kwargs.pop("background", False)
-        df = super().collect(*args, background=background, **kwargs)
+        df: pl.DataFrame = super().collect(*args, background=background, **kwargs)
+        df = DataFrame(df)
         if getattr(self, "model", False):
-            cls = DataFrame._construct_dataframe_model_class(model=self.model)
-        else:
-            cls = DataFrame
-        return cls._from_pydf(df._df)
+            df = df.set_model(self.model)
+        return df
 
     def derive(self: LDF, columns: list[str] | None = None) -> LDF:
         """Populate columns which have ``pt.Field(derived_from=...)`` definitions.
@@ -148,33 +188,35 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
 
     def _derive_column(
         self,
-        df: LDF,
+        lf: LDF,
         column_name: str,
-        column_infos: Dict[str, ColumnInfo],
-    ) -> Tuple[LDF, Sequence[str]]:
+        column_infos: dict[str, ColumnInfo],
+    ) -> tuple[LDF, Sequence[str]]:
         if (
             column_infos.get(column_name, None) is None
             or column_infos[column_name].derived_from is None
         ):
-            return df, []
+            return lf, []
+
         derived_from = column_infos[column_name].derived_from
         dtype = self.model.dtypes[column_name]
         derived_columns = []
+
         if isinstance(derived_from, str):
-            df = df.with_columns(pl.col(derived_from).cast(dtype).alias(column_name))
+            lf = lf.with_columns(pl.col(derived_from).cast(dtype).alias(column_name))
         elif isinstance(derived_from, pl.Expr):
             root_cols = derived_from.meta.root_names()
             while root_cols:
                 root_col = root_cols.pop()
-                df, _derived_columns = self._derive_column(df, root_col, column_infos)
+                lf, _derived_columns = self._derive_column(lf, root_col, column_infos)
                 derived_columns.extend(_derived_columns)
-            df = df.with_columns(derived_from.cast(dtype).alias(column_name))
+            lf = lf.with_columns(derived_from.cast(dtype).alias(column_name))
         else:
             raise TypeError(
                 "Can not derive dataframe column from type " f"{type(derived_from)}."
             )
         derived_columns.append(column_name)
-        return df, derived_columns
+        return lf, derived_columns
 
     def unalias(self: LDF) -> LDF:
         """Un-aliases column names using information from pydantic validation_alias.
@@ -191,7 +233,7 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
             return self
         exprs = []
 
-        def to_expr(va: str | AliasPath | AliasChoices) -> Optional[pl.Expr]:
+        def to_expr(va: str | AliasPath | AliasChoices) -> pl.Expr | None:
             if isinstance(va, str):
                 return pl.col(va) if va in self.collect_schema() else None
             elif isinstance(va, AliasPath):
@@ -200,12 +242,12 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
                         f"TODO figure out how this AliasPath behaves ({va})"
                     )
                 return (
-                    pl.col(va.path[0]).list.get(va.path[1], null_on_oob=True)
+                    pl.col(str(va.path[0])).list.get(va.path[1], null_on_oob=True)
                     if va.path[0] in self.collect_schema()
                     else None
                 )
             elif isinstance(va, AliasChoices):
-                local_expr: Optional[pl.Expr] = None
+                local_expr: pl.Expr | None = None
                 for choice in va.choices:
                     if (part := to_expr(choice)) is not None:
                         local_expr = (
@@ -235,7 +277,7 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
         return self.select(exprs)
 
     def cast(
-        self: LDF, strict: bool = False, columns: Optional[Sequence[str]] = None
+        self: LDF, strict: bool = False, columns: Sequence[str] | None = None
     ) -> LDF:
         """Cast columns to `dtypes` specified by the associated Patito model.
 
@@ -292,9 +334,12 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
         return self.with_columns(exprs)
 
     @classmethod
-    def from_existing(cls: Type[LDF], lf: pl.LazyFrame) -> LDF:
+    def from_existing(cls: type[LDF], lf: pl.LazyFrame) -> LDF:
         """Construct a patito.DataFrame object from an existing polars.DataFrame object."""
-        return cls.model.LazyFrame._from_pyldf(lf._ldf).cast()
+        if getattr(cls, "model", False):
+            return cls.model.LazyFrame._from_pyldf(super().lazy()._ldf)  # type: ignore
+
+        return LazyFrame._from_pyldf(lf._ldf)  # type: ignore
 
 
 class DataFrame(pl.DataFrame, Generic[ModelType]):
@@ -326,31 +371,7 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
     :ref:`Product.validate <DataFrame.validate>`.
     """
 
-    model: Type[ModelType]
-
-    @classmethod
-    def _construct_dataframe_model_class(
-        cls: Type[DF], model: Type[OtherModelType]
-    ) -> Type[DataFrame[OtherModelType]]:
-        """Return custom DataFrame sub-class where DataFrame.model is set.
-
-        Can be used to construct a DataFrame class where
-        DataFrame.set_model(model) is implicitly invoked at instantiation.
-
-        Args:
-            model: A patito model which should be used to validate the dataframe.
-
-        Returns:
-            A custom DataFrame model class where DataFrame._model has been correctly
-                "hard-coded" to the given model.
-
-        """
-        new_class = type(
-            f"{model.model_json_schema()['title']}DataFrame",
-            (cls,),
-            {"model": model},
-        )
-        return new_class
+    model: type[ModelType]
 
     def lazy(self: DataFrame[ModelType]) -> LazyFrame[ModelType]:
         """Convert DataFrame into LazyFrame.
@@ -361,15 +382,12 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
             A new LazyFrame object.
 
         """
-        lazyframe_class: LazyFrame[ModelType] = (
-            LazyFrame._construct_lazyframe_model_class(
-                model=getattr(self, "model", None)
-            )
-        )  # type: ignore
-        ldf = lazyframe_class._from_pyldf(super().lazy()._ldf)
-        return ldf
+        if getattr(self, "model", False):
+            return self.model.LazyFrame._from_pyldf(super().lazy()._ldf)  # type: ignore
+
+        return LazyFrame._from_pyldf(super().lazy()._ldf)  # type: ignore
 
-    def set_model(self, model):  # type: ignore[no-untyped-def] # noqa: ANN001, ANN201
+    def set_model(self, model: type[OtherModelType]) -> DataFrame[OtherModelType]:
         """Associate a given patito ``Model`` with the dataframe.
 
         The model schema is used by methods that depend on a model being associated with
@@ -425,11 +443,7 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
             >>> casted_classes.validate()
 
         """
-        cls = self._construct_dataframe_model_class(model=model)
-        return cast(
-            DataFrame[model],
-            cls._from_pydf(self._df),
-        )
+        return model.DataFrame(self._df)
 
     def unalias(self: DF) -> DF:
         """Un-aliases column names using information from pydantic validation_alias.
@@ -445,7 +459,7 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
         return self.lazy().unalias().collect()
 
     def cast(
-        self: DF, strict: bool = False, columns: Optional[Sequence[str]] = None
+        self: DF, strict: bool = False, columns: Sequence[str] | None = None
     ) -> DF:
         """Cast columns to `dtypes` specified by the associated Patito model.
 
@@ -489,8 +503,7 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
 
     def drop(
         self: DF,
-        columns: Optional[Union[str, Collection[str]]] = None,
-        *more_columns: str,
+        columns: str | Collection[str] | None = None,
     ) -> DF:
         """Drop one or more columns from the dataframe.
 
@@ -502,7 +515,6 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
             columns: A single column string name, or list of strings, indicating
                 which columns to drop. If not specified, all columns *not*
                 specified by the associated dataframe model will be dropped.
-            more_columns: Additional named columns to drop.
 
         Returns:
             DataFrame[Model]: New dataframe without the specified columns.
@@ -525,27 +537,29 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
 
         """
         if columns is not None:
-            return self._from_pydf(super().drop(columns)._df)
+            # I get a single null row if I try to use super() here, so go via
+            # pl.DataFrame instead.
+            return self._from_pydf(pl.DataFrame(self._df).drop(columns)._df)
         else:
             return self.drop(list(set(self.columns) - set(self.model.columns)))
 
-    def validate(self, columns: Optional[Sequence[str]] = None, **kwargs: Any):
+    def validate(self, columns: Sequence[str] | None = None, **kwargs: Any):
         """Validate the schema and content of the dataframe.
 
         You must invoke ``.set_model()`` before invoking ``.validate()`` in order
         to specify how the dataframe should be validated.
 
         Returns:
-            DataFrame[Model]: The original dataframe, if correctly validated.
+            DataFrame[Model]: The original patito dataframe, if correctly validated.
 
         Raises:
+            patito.exceptions.DataFrameValidationError: If the dataframe does not match the
+                specified schema.
+
             TypeError: If ``DataFrame.set_model()`` has not been invoked prior to
                 validation. Note that ``patito.Model.DataFrame`` automatically invokes
                 ``DataFrame.set_model()`` for you.
 
-            patito.exceptions.DataFrameValidationError: If the dataframe does not match the
-                specified schema.
-
         Examples:
             >>> import patito as pt
 
@@ -623,13 +637,12 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
 
     def fill_null(
         self: DF,
-        value: Optional[Any] = None,
-        strategy: Optional[
-            Literal[
-                "forward", "backward", "min", "max", "mean", "zero", "one", "defaults"
-            ]
-        ] = None,
-        limit: Optional[int] = None,
+        value: Any | None = None,
+        strategy: Literal[
+            "forward", "backward", "min", "max", "mean", "zero", "one", "defaults"
+        ]
+        | None = None,
+        limit: int | None = None,
         matches_supertype: bool = True,
     ) -> DF:
         """Fill null values using a filling strategy, literal, or ``Expr``.
@@ -689,14 +702,13 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
                         pl.lit(default_value, self.model.dtypes[column])
                     )
                     if column in self.columns
-                    else pl.Series(column, [default_value], self.model.dtypes[column])
-                )  # NOTE: hack to get around polars bug https://github.com/pola-rs/polars/issues/13602
-                # else pl.lit(default_value, self.model.dtypes[column]).alias(column)
+                    else pl.lit(default_value, self.model.dtypes[column]).alias(column)
+                )
                 for column, default_value in self.model.defaults.items()
             ]
-        ).set_model(self.model)
+        ).set_model(self.model)  # type: ignore
 
-    def get(self, predicate: Optional[pl.Expr] = None) -> ModelType:
+    def get(self, predicate: pl.Expr | None = None) -> ModelType:
         """Fetch the single row that matches the given polars predicate.
 
         If you expect a data frame to already consist of one single row,
@@ -778,7 +790,70 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
         else:
             return self._pydantic_model().from_row(row)  # type: ignore
 
-    def _pydantic_model(self) -> Type[Model]:
+    def iter_models(
+        self, validate_df: bool = True, validate_model: bool = False
+    ) -> ModelGenerator[ModelType]:
+        """Iterate over all rows in the dataframe as pydantic models.
+
+        Args:
+            validate_df: If set to ``True``, the dataframe will be validated before
+                making models out of each row. If set to ``False``, beware that columns
+                need to be the exact same as the model fields.
+            validate_model: If set to ``True``, each model will be validated when
+                constructing. Disabled by default since df validation should cover this case.
+
+        Yields:
+            Model: A pydantic-derived model representing the given row. .to_list() can be
+                used to convert the iterator to a list.
+
+        Raises:
+            TypeError: If ``DataFrame.set_model()`` has not been invoked prior to
+                iteration.
+
+        Example:
+            >>> import patito as pt
+            >>> import polars as pl
+            >>> class Product(pt.Model):
+            ...     product_id: int = pt.Field(unique=True)
+            ...     price: float
+
+            >>> df = pt.DataFrame({"product_id": [1, 2], "price": [10., 20.]})
+            >>> df = df.set_model(Product)
+            >>> for product in df.iter_models():
+            ...     print(product)
+            ...
+            Product(product_id=1, price=10.0)
+            Product(product_id=2, price=20.0)
+
+        """
+        if not hasattr(self, "model"):
+            raise TypeError(
+                f"You must invoke {self.__class__.__name__}.set_model() "
+                f"before invoking {self.__class__.__name__}.iter_models()."
+            )
+
+        df = self.validate(drop_superfluous_columns=True) if validate_df else self
+
+        def _iter_models_with_validate(
+            _df: DataFrame[ModelType],
+        ) -> Iterator[ModelType]:
+            for row in _df.iter_rows(named=True):
+                yield self.model(**row)
+
+        def _iter_models_without_validate(
+            _df: DataFrame[ModelType],
+        ) -> Iterator[ModelType]:
+            for row in _df.iter_rows(named=True):
+                yield self.model.model_construct(**row)
+
+        _iter_models = (
+            _iter_models_with_validate
+            if validate_model
+            else _iter_models_without_validate
+        )
+        return ModelGenerator(_iter_models(df))
+
+    def _pydantic_model(self) -> type[Model]:
         """Dynamically construct patito model compliant with dataframe.
 
         Returns:
@@ -790,7 +865,7 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
 
         pydantic_annotations = {column: (Any, ...) for column in self.columns}
         return cast(
-            Type[Model],
+            type[Model],
             create_model(  # type: ignore
                 "UntypedRow",
                 __base__=Model,
@@ -804,7 +879,7 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
 
     @classmethod
     def read_csv(  # type: ignore[no-untyped-def]
-        cls: Type[DF],
+        cls: type[DF],
         *args,  # noqa: ANN002
         **kwargs,  # noqa: ANN003
     ) -> DF:
@@ -888,15 +963,13 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
     # --- Type annotation overrides ---
     def filter(  # noqa: D102
         self: DF,
-        predicate: Union[
-            pl.Expr, str, pl.Series, list[bool], np.ndarray[Any, Any], bool
-        ],
+        predicate: pl.Expr | str | pl.Series | list[bool] | np.ndarray[Any, Any] | bool,
     ) -> DF:
         return cast(DF, super().filter(predicate))
 
     def select(  # noqa: D102
         self: DF,
-        *exprs: Union[IntoExpr, Iterable[IntoExpr]],
+        *exprs: IntoExpr | Iterable[IntoExpr],
         **named_exprs: IntoExpr,
     ) -> DF:
         return cast(  # pyright: ignore[redundant-cast]
@@ -905,7 +978,7 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
 
     def with_columns(  # noqa: D102
         self: DF,
-        *exprs: Union[IntoExpr, Iterable[IntoExpr]],
+        *exprs: IntoExpr | Iterable[IntoExpr],
         **named_exprs: IntoExpr,
     ) -> DF:
         return cast(DF, super().with_columns(*exprs, **named_exprs))