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

patito/polars.py

@@ -9,6 +9,7 @@ from typing import (
     Dict,
     Generic,
     Iterable,
+    Literal,
     Optional,
     Sequence,
     Tuple,
@@ -19,9 +20,8 @@ from typing import (
 )
 
 import polars as pl
-from polars.type_aliases import IntoExpr
+from polars._typing import IntoExpr
 from pydantic import AliasChoices, AliasPath, create_model
-from typing_extensions import Literal
 
 from patito._pydantic.column_info import ColumnInfo
 from patito.exceptions import MultipleRowsReturned, RowDoesNotExist
@@ -53,12 +53,10 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
         DataFrame.set_model(model) is implicitly invoked at collection.
 
         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.
 
         Returns:
-        -------
             A custom LazyFrame model class where LazyFrame.model has been correctly
                 "hard-coded" to the given model.
 
@@ -101,21 +99,17 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
         result of which will be used to populate the column values.
 
         Args:
-        ----
             columns: Optionally, a list of column names to derive. If not provided, all
                 columns are used.
 
         Returns:
-        -------
             DataFrame[Model]: A new dataframe where all derivable columns are provided.
 
         Raises:
-        ------
             TypeError: If the ``derived_from`` parameter of ``patito.Field`` is given
                 as something else than a string or polars expression.
 
         Examples:
-        --------
             >>> import patito as pt
             >>> import polars as pl
             >>> class Foo(pt.Model):
@@ -136,7 +130,7 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
         """
         derived_columns = []
         props = self.model._schema_properties()
-        original_columns = set(self.columns)
+        original_columns = set(self.collect_schema())
         to_derive = self.model.derived_columns if columns is None else columns
         for column_name in to_derive:
             if column_name not in derived_columns:
@@ -189,8 +183,7 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
 
         limitation - AliasChoice validation type only supports selecting a single element of an array
 
-        Returns
-        -------
+        Returns:
             DataFrame[Model]: A dataframe with columns normalized to model names.
 
         """
@@ -200,15 +193,15 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
 
         def to_expr(va: str | AliasPath | AliasChoices) -> Optional[pl.Expr]:
             if isinstance(va, str):
-                return pl.col(va) if va in self.columns else None
+                return pl.col(va) if va in self.collect_schema() else None
             elif isinstance(va, AliasPath):
                 if len(va.path) != 2 or not isinstance(va.path[1], int):
                     raise NotImplementedError(
                         f"TODO figure out how this AliasPath behaves ({va})"
                     )
                 return (
-                    pl.col(va.path[0]).list.get(va.path[1])
-                    if va.path[0] in self.columns
+                    pl.col(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):
@@ -231,7 +224,7 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
                 exprs.append(pl.col(name))
             else:
                 expr = to_expr(field_info.validation_alias)
-                if name in self.columns:
+                if name in self.collect_schema().names():
                     if expr is None:
                         exprs.append(pl.col(name))
                     else:
@@ -247,7 +240,6 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
         """Cast columns to `dtypes` specified by the associated Patito model.
 
         Args:
-        ----
             strict: If set to ``False``, columns which are technically compliant with
                 the specified field type, will not be casted. For example, a column
                 annotated with ``int`` is technically compliant with ``pl.UInt8``, even
@@ -258,11 +250,9 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
                 columns are casted.
 
         Returns:
-        -------
             LazyFrame[Model]: A dataframe with columns casted to the correct dtypes.
 
         Examples:
-        --------
             Create a simple model:
 
             >>> import patito as pt
@@ -288,9 +278,9 @@ class LazyFrame(pl.LazyFrame, Generic[ModelType]):
         properties = self.model._schema_properties()
         valid_dtypes = self.model.valid_dtypes
         default_dtypes = self.model.dtypes
-        columns = columns or self.columns
+        columns = columns or self.collect_schema().names()
         exprs = []
-        for column, current_dtype in zip(self.columns, self.dtypes):
+        for column, current_dtype in self.collect_schema().items():
             if (column not in columns) or (column not in properties):
                 exprs.append(pl.col(column))
             elif "dtype" in properties[column]:
@@ -348,11 +338,9 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
         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.
 
@@ -369,15 +357,14 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
 
         See documentation of polars.DataFrame.lazy() for full description.
 
-        Returns
-        -------
+        Returns:
             A new LazyFrame object.
 
         """
-        lazyframe_class: LazyFrame[
-            ModelType
-        ] = LazyFrame._construct_lazyframe_model_class(
-            model=getattr(self, "model", None)
+        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
@@ -392,17 +379,14 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
         ``DataFrame(...).set_model(Model)`` is equivalent with ``Model.DataFrame(...)``.
 
         Args:
-        ----
             model (Model): Sub-class of ``patito.Model`` declaring the schema of the
                 dataframe.
 
         Returns:
-        -------
             DataFrame[Model]: Returns the same dataframe, but with an attached model
             that is required for certain model-specific dataframe methods to work.
 
         Examples:
-        --------
             >>> from typing_extensions import Literal
             >>> import patito as pt
             >>> import polars as pl
@@ -454,8 +438,7 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
 
         limitation - AliasChoice validation type only supports selecting a single element of an array
 
-        Returns
-        -------
+        Returns:
             DataFrame[Model]: A dataframe with columns normalized to model names.
 
         """
@@ -467,7 +450,6 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
         """Cast columns to `dtypes` specified by the associated Patito model.
 
         Args:
-        ----
             strict: If set to ``False``, columns which are technically compliant with
                 the specified field type, will not be casted. For example, a column
                 annotated with ``int`` is technically compliant with ``pl.UInt8``, even
@@ -478,11 +460,9 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
                 columns are casted.
 
         Returns:
-        -------
             DataFrame[Model]: A dataframe with columns casted to the correct dtypes.
 
         Examples:
-        --------
             Create a simple model:
 
             >>> import patito as pt
@@ -519,18 +499,15 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
         :ref:`DataFrame.set_model <DataFrame.set_model>`, are dropped.
 
         Args:
-        ----
             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.
 
         Examples:
-        --------
             >>> import patito as pt
             >>> class Model(pt.Model):
             ...     column_1: int
@@ -552,20 +529,16 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
         else:
             return self.drop(list(set(self.columns) - set(self.model.columns)))
 
-    def validate(
-        self: DF, columns: Optional[Sequence[str]] = None, **kwargs: Any
-    ) -> DF:
+    def validate(self, columns: Optional[Sequence[str]] = 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
-        -------
+        Returns:
             DataFrame[Model]: The original dataframe, if correctly validated.
 
-        Raises
-        ------
+        Raises:
             TypeError: If ``DataFrame.set_model()`` has not been invoked prior to
                 validation. Note that ``patito.Model.DataFrame`` automatically invokes
                 ``DataFrame.set_model()`` for you.
@@ -573,8 +546,7 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
             patito.exceptions.DataFrameValidationError: If the dataframe does not match the
                 specified schema.
 
-        Examples
-        --------
+        Examples:
             >>> import patito as pt
 
 
@@ -621,17 +593,14 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
         column name. Alternatively, an arbitrary polars expression can be given, the
         result of which will be used to populate the column values.
 
-        Returns
-        -------
+        Returns:
             DataFrame[Model]: A new dataframe where all derivable columns are provided.
 
-        Raises
-        ------
+        Raises:
             TypeError: If the ``derived_from`` parameter of ``patito.Field`` is given
                 as something else than a string or polars expression.
 
-        Examples
-        --------
+        Examples:
             >>> import patito as pt
             >>> import polars as pl
             >>> class Foo(pt.Model):
@@ -665,11 +634,10 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
     ) -> DF:
         """Fill null values using a filling strategy, literal, or ``Expr``.
 
-        If ``"default"`` is provided as the strategy, the model fields with default
+        If ``"defaults"`` is provided as the strategy, the model fields with default
         values are used to fill missing values.
 
         Args:
-        ----
             value: Value used to fill null values.
             strategy: Accepts the same arguments as ``polars.DataFrame.fill_null`` in
                 addition to ``"defaults"`` which will use the field's default value if
@@ -680,12 +648,10 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
 
 
         Returns:
-        -------
             DataFrame[Model]: A new dataframe with nulls filled in according to the
             provided ``strategy`` parameter.
 
         Example:
-        -------
             >>> import patito as pt
             >>> class Product(pt.Model):
             ...     name: str
@@ -737,7 +703,6 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
         you can use ``.get()`` without any arguments to return that row.
 
         Raises:
-        ------
             RowDoesNotExist: If zero rows evaluate to true for the given predicate.
             MultipleRowsReturned: If more than one row evaluates to true for the given
                 predicate.
@@ -746,15 +711,12 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
                 same class.
 
         Args:
-        ----
             predicate: A polars expression defining the criteria of the filter.
 
         Returns:
-        -------
             Model: A pydantic-derived base model representing the given row.
 
         Example:
-        -------
             >>> import patito as pt
             >>> import polars as pl
             >>> df = pt.DataFrame({"product_id": [1, 2, 3], "price": [10, 10, 20]})
@@ -819,8 +781,7 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
     def _pydantic_model(self) -> Type[Model]:
         """Dynamically construct patito model compliant with dataframe.
 
-        Returns
-        -------
+        Returns:
             A pydantic model class where all the rows have been specified as
                 `typing.Any` fields.
 
@@ -853,16 +814,13 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
         to populate the given column(s).
 
         Args:
-        ----
             *args: All positional arguments are forwarded to ``polars.read_csv``.
             **kwargs: All keyword arguments are forwarded to ``polars.read_csv``.
 
         Returns:
-        -------
             DataFrame[Model]: A dataframe representing the given CSV file data.
 
         Examples:
-        --------
             The ``DataFrame.read_csv`` method can be used to automatically set the
             correct column names when reading CSV files without headers.
 
@@ -907,9 +865,23 @@ class DataFrame(pl.DataFrame, Generic[ModelType]):
             # └─────┴─────┘
 
         """
-        kwargs.setdefault("dtypes", cls.model.dtypes)
-        if not kwargs.get("has_header", True) and "columns" not in kwargs:
+        kwargs.setdefault("schema_overrides", cls.model.dtypes)
+        has_header = kwargs.get("has_header", True)
+        if not has_header and "columns" not in kwargs:
             kwargs.setdefault("new_columns", cls.model.columns)
+        alias_gen = cls.model.model_config.get("alias_generator")
+        if alias_gen:
+            alias_func = alias_gen.validation_alias or alias_gen.alias
+        if has_header and alias_gen and alias_func:
+            fields_to_cols = {
+                field_name: alias_func(field_name)
+                for field_name in cls.model.model_fields
+            }
+            kwargs["schema_overrides"] = {
+                fields_to_cols.get(field, field): dtype
+                for field, dtype in kwargs["schema_overrides"].items()
+            }
+            # TODO: other forms of alias setting like in Field
         df = cls.model.DataFrame._from_pydf(pl.read_csv(*args, **kwargs)._df)
         return df.derive()