dataframely 1.12.0__tar.gz → 1.13.0__tar.gz

{dataframely-1.12.0 → dataframely-1.13.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataframely
-Version: 1.12.0
+Version: 1.13.0
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
@@ -12,9 +12,11 @@ Requires-Dist: typing-extensions ; python_full_version < '3.11'
 Requires-Dist: deltalake ; extra == 'deltalake'
 Requires-Dist: sqlalchemy ; extra == 'sqlalchemy'
 Requires-Dist: pyarrow ; extra == 'pyarrow'
+Requires-Dist: pydantic>=2 ; extra == 'pydantic'
 Provides-Extra: deltalake
 Provides-Extra: sqlalchemy
 Provides-Extra: pyarrow
+Provides-Extra: pydantic
 License-File: LICENSE
 Summary: A declarative, polars-native data frame validation library
 Author-email: Andreas Albert <andreas.albert@quantco.com>, Daniel Elsner <daniel.elsner@quantco.com>, Oliver Borchert <oliver.borchert@quantco.com>

{dataframely-1.12.0 → dataframely-1.13.0}/dataframely/_base_schema.py

@@ -5,10 +5,10 @@ from __future__ import annotations
 
 import sys
 import textwrap
-from abc import ABCMeta
+from abc import ABCMeta, abstractmethod
 from copy import copy
 from dataclasses import dataclass, field
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import polars as pl
 
@@ -21,6 +21,10 @@ if sys.version_info >= (3, 11):
 else:
     from typing_extensions import Self
 
+
+if TYPE_CHECKING:
+    from ._typing import DataFrame
+
 _COLUMN_ATTR = "__dataframely_columns__"
 _RULE_ATTR = "__dataframely_rules__"
 
@@ -198,11 +202,46 @@ class BaseSchema(metaclass=SchemaMeta):
             columns[name]._name = name
         return columns
 
+    @classmethod
+    @abstractmethod
+    def polars_schema(cls) -> pl.Schema:
+        """Obtain the polars schema for this schema.
+
+        Returns:
+            A :mod:`polars` schema that mirrors the schema defined by this class.
+        """
+
     @classmethod
     def primary_keys(cls) -> list[str]:
         """The primary key columns in this schema (possibly empty)."""
         return _primary_keys(cls.columns())
 
+    @classmethod
+    @abstractmethod
+    def validate(
+        cls, df: pl.DataFrame | pl.LazyFrame, /, *, cast: bool = False
+    ) -> DataFrame[Self]:
+        """Validate that a data frame satisfies the schema.
+
+        Args:
+            df: The data frame to validate.
+            cast: Whether columns with a wrong data type in the input data frame are
+                cast to the schema's defined data type if possible.
+
+        Returns:
+            The (collected) input data frame, wrapped in a generic version of the
+            input's data frame type to reflect schema adherence. The data frame is
+            guaranteed to maintain its order.
+
+        Raises:
+            ValidationError: If the input data frame does not satisfy the schema
+                definition.
+
+        Note:
+            This method _always_ collects the input data frame in order to raise
+            potential validation errors.
+        """
+
     @classmethod
     def _validation_rules(cls, *, with_cast: bool) -> dict[str, Rule]:
         return _build_rules(

{dataframely-1.12.0 → dataframely-1.13.0}/dataframely/_compat.py

@@ -54,6 +54,19 @@ try:
 except ImportError:  # pragma: no cover
     pa = _DummyModule("pyarrow")
 
+
+# -------------------------------------- PYDANTIC ------------------------------------ #
+
+try:
+    import pydantic
+except ImportError:  # pragma: no cover
+    pydantic = _DummyModule("pydantic")  # type: ignore
+
+try:
+    from pydantic_core import core_schema as pydantic_core_schema  # pragma: no cover
+except ImportError:
+    pydantic_core_schema = _DummyModule("pydantic_core_schema")  # type: ignore
+
 # ------------------------------------------------------------------------------------ #
 
 __all__ = [
@@ -64,4 +77,6 @@ __all__ = [
     "pa",
     "MSDialect_pyodbc",
     "PGDialect_psycopg2",
+    "pydantic",
+    "pydantic_core_schema",
 ]

dataframely-1.13.0/dataframely/_pydantic.py

@@ -0,0 +1,115 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+from __future__ import annotations
+
+from functools import partial
+from typing import TYPE_CHECKING, Literal, TypeVar, get_args, get_origin, overload
+
+import polars as pl
+
+from ._base_schema import BaseSchema
+from ._compat import pydantic, pydantic_core_schema
+from .exc import ValidationError
+
+if TYPE_CHECKING:
+    from ._typing import DataFrame, LazyFrame
+
+
+_S = TypeVar("_S", bound=BaseSchema)
+
+
+def _dict_to_df(schema_type: type[BaseSchema], data: dict) -> pl.DataFrame:
+    return pl.from_dict(
+        data,
+        schema=schema_type.polars_schema(),
+    )
+
+
+def _validate_df_schema(schema_type: type[_S], df: pl.DataFrame) -> DataFrame[_S]:
+    try:
+        return schema_type.validate(df, cast=False)
+    except ValidationError as e:
+        raise ValueError("DataFrame violates schema") from e
+
+
+def _serialize_df(df: pl.DataFrame) -> dict:
+    return df.to_dict(as_series=False)
+
+
+@overload
+def get_pydantic_core_schema(
+    source_type: type[DataFrame],
+    _handler: pydantic.GetCoreSchemaHandler,
+    lazy: Literal[False],
+) -> pydantic_core_schema.CoreSchema: ...
+
+
+@overload
+def get_pydantic_core_schema(
+    source_type: type[LazyFrame],
+    _handler: pydantic.GetCoreSchemaHandler,
+    lazy: Literal[True],
+) -> pydantic_core_schema.CoreSchema: ...
+
+
+def get_pydantic_core_schema(
+    source_type: type[DataFrame | LazyFrame],
+    _handler: pydantic.GetCoreSchemaHandler,
+    lazy: bool,
+) -> pydantic_core_schema.CoreSchema:
+    # https://docs.pydantic.dev/2.11/concepts/types/#handling-custom-generic-classes
+    origin = get_origin(source_type)
+    if origin is None:
+        # used as `x: dy.DataFrame` without schema
+        raise TypeError("DataFrame must be parametrized with a schema")
+
+    schema_type: type[BaseSchema] = get_args(source_type)[0]
+
+    # accept a DataFrame, a LazyFrame, or a dict that is converted to a DataFrame
+    # (-> output: DataFrame or LazyFrame)
+    polars_schema = pydantic_core_schema.union_schema(
+        [
+            pydantic_core_schema.is_instance_schema(pl.DataFrame),
+            pydantic_core_schema.is_instance_schema(pl.LazyFrame),
+            pydantic_core_schema.chain_schema(
+                [
+                    pydantic_core_schema.dict_schema(),
+                    pydantic_core_schema.no_info_plain_validator_function(
+                        partial(_dict_to_df, schema_type)
+                    ),
+                ]
+            ),
+        ]
+    )
+
+    to_lazy_schema = []
+    if lazy:
+        # If the Pydantic field type is LazyFrame, add a step to convert
+        # the model back to a LazyFrame.
+        to_lazy_schema.append(
+            pydantic_core_schema.no_info_plain_validator_function(
+                lambda df: df.lazy(),
+            )
+        )
+
+    return pydantic_core_schema.chain_schema(
+        [
+            polars_schema,
+            pydantic_core_schema.no_info_plain_validator_function(
+                partial(_validate_df_schema, schema_type)
+            ),
+            *to_lazy_schema,
+        ],
+        serialization=pydantic_core_schema.plain_serializer_function_ser_schema(
+            _serialize_df
+        ),
+    )
+
+
+def get_pydantic_json_schema(
+    handler: pydantic.GetJsonSchemaHandler,
+) -> pydantic.json_schema.JsonSchemaValue:
+    from pydantic_core import core_schema
+
+    # This could be made more sophisticated by actually reflecting the schema.
+    return handler(core_schema.dict_schema())

{dataframely-1.12.0 → dataframely-1.13.0}/dataframely/_typing.py

@@ -9,6 +9,8 @@ from typing import TYPE_CHECKING, Any, Concatenate, Generic, Literal, ParamSpec,
 import polars as pl
 
 from ._base_schema import BaseSchema
+from ._compat import pydantic, pydantic_core_schema
+from ._pydantic import get_pydantic_core_schema, get_pydantic_json_schema
 
 S = TypeVar("S", bound=BaseSchema, covariant=True)
 
@@ -70,6 +72,20 @@ class DataFrame(pl.DataFrame, Generic[S]):
     def shrink_to_fit(self, *args: Any, **kwargs: Any) -> DataFrame[S]:
         raise NotImplementedError  # pragma: no cover
 
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
+    ) -> pydantic_core_schema.CoreSchema:
+        return get_pydantic_core_schema(source_type, handler, lazy=False)
+
+    @classmethod
+    def __get_pydantic_json_schema__(
+        cls,
+        _core_schema: pydantic_core_schema.CoreSchema,
+        handler: pydantic.GetJsonSchemaHandler,
+    ) -> pydantic.json_schema.JsonSchemaValue:
+        return get_pydantic_json_schema(handler)
+
 
 class LazyFrame(pl.LazyFrame, Generic[S]):
     """Generic wrapper around a :class:`polars.LazyFrame` to attach schema information.
@@ -113,3 +129,17 @@ class LazyFrame(pl.LazyFrame, Generic[S]):
     @inherit_signature(pl.LazyFrame.set_sorted)
     def set_sorted(self, *args: Any, **kwargs: Any) -> LazyFrame[S]:
         raise NotImplementedError  # pragma: no cover
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
+    ) -> pydantic_core_schema.CoreSchema:
+        return get_pydantic_core_schema(source_type, handler, lazy=True)
+
+    @classmethod
+    def __get_pydantic_json_schema__(
+        cls,
+        _core_schema: pydantic_core_schema.CoreSchema,
+        handler: pydantic.GetJsonSchemaHandler,
+    ) -> pydantic.json_schema.JsonSchemaValue:
+        return get_pydantic_json_schema(handler)

{dataframely-1.12.0 → dataframely-1.13.0}/dataframely/columns/_base.py

@@ -381,15 +381,6 @@ class Column(ABC):
         if name == "check":
             return _compare_checks(lhs, rhs, column_expr)
 
-        lhs_is_series = isinstance(lhs, pl.Series)
-        rhs_is_series = isinstance(rhs, pl.Series)
-
-        if lhs_is_series != rhs_is_series:
-            return False
-
-        if lhs_is_series and rhs_is_series:
-            return _compare_series(lhs, rhs)
-
         return lhs == rhs
 
     # -------------------------------- DUNDER METHODS -------------------------------- #
@@ -413,10 +404,6 @@ class Column(ABC):
         return self.__class__.__name__.lower()
 
 
-def _compare_series(lhs: pl.Series, rhs: pl.Series) -> bool:
-    return (len(lhs) == len(rhs)) and lhs.equals(rhs)
-
-
 def _compare_checks(lhs: Check | None, rhs: Check | None, expr: pl.Expr) -> bool:
     match (lhs, rhs):
         case (None, None):

{dataframely-1.12.0 → dataframely-1.13.0}/dataframely/columns/enum.py

@@ -67,12 +67,8 @@ class Enum(Column):
             metadata=metadata,
         )
         if isclass(categories) and issubclass(categories, enum.Enum):
-            categories = pl.Series(
-                values=[getattr(v, "value", v) for v in categories.__members__.values()]
-            )
-        elif not isinstance(categories, pl.Series):
-            categories = pl.Series(values=categories)
-        self.categories = categories
+            categories = (item.value for item in categories)
+        self.categories = list(categories)
 
     @property
     def dtype(self) -> pl.DataType:
@@ -81,7 +77,7 @@ class Enum(Column):
     def validate_dtype(self, dtype: PolarsDataType) -> bool:
         if not isinstance(dtype, pl.Enum):
             return False
-        return self.categories.equals(dtype.categories)
+        return self.categories == dtype.categories.to_list()
 
     def sqlalchemy_dtype(self, dialect: sa.Dialect) -> sa_TypeEngine:
         category_lengths = [len(c) for c in self.categories]
@@ -102,6 +98,6 @@ class Enum(Column):
     def _sample_unchecked(self, generator: Generator, n: int) -> pl.Series:
         return generator.sample_choice(
             n,
-            choices=self.categories.to_list(),
+            choices=self.categories,
             null_probability=self._null_probability,
         ).cast(self.dtype)

{dataframely-1.12.0 → dataframely-1.13.0}/dataframely/schema.py

@@ -414,26 +414,6 @@ class Schema(BaseSchema, ABC):
     def validate(
         cls, df: pl.DataFrame | pl.LazyFrame, /, *, cast: bool = False
     ) -> DataFrame[Self]:
-        """Validate that a data frame satisfies the schema.
-
-        Args:
-            df: The data frame to validate.
-            cast: Whether columns with a wrong data type in the input data frame are
-                cast to the schema's defined data type if possible.
-
-        Returns:
-            The (collected) input data frame, wrapped in a generic version of the
-            input's data frame type to reflect schema adherence. The data frame is
-            guaranteed to maintain its order.
-
-        Raises:
-            ValidationError: If the input data frame does not satisfy the schema
-                definition.
-
-        Note:
-            This method _always_ collects the input data frame in order to raise
-            potential validation errors.
-        """
         # We can dispatch to the `filter` method and raise an error if any row cannot
         # be validated
         df_valid, failures = cls.filter(df, cast=cast)
@@ -1118,11 +1098,6 @@ class Schema(BaseSchema, ABC):
 
     @classmethod
     def polars_schema(cls) -> pl.Schema:
-        """Obtain the polars schema for this schema.
-
-        Returns:
-            A :mod:`polars` schema that mirrors the schema defined by this class.
-        """
         return pl.Schema({name: col.dtype for name, col in cls.columns().items()})
 
     @classmethod

{dataframely-1.12.0 → dataframely-1.13.0}/docs/sites/quickstart.rst

@@ -253,6 +253,7 @@ Lastly, ``dataframely`` schemas can be used to integrate with external tools:
 - ``HouseSchema.create_empty()`` creates an empty ``dy.DataFrame[HouseSchema]`` that can be used for testing
 - ``HouseSchema.sql_schema()`` provides a list of `sqlalchemy <https://www.sqlalchemy.org>`_ columns that can be used to create SQL tables using types and constraints in line with the schema
 - ``HouseSchema.pyarrow_schema()`` provides a `pyarrow <https://arrow.apache.org/docs/python/index.html>`_ schema with appropriate column dtypes and nullability information
+- You can use ``dy.DataFrame[HouseSchema]`` (or the ``LazyFrame`` equivalent) as fields in `pydantic <https://pydantic.dev>`_ models, including support for validation and serialization. Integration with pydantic is unstable.
 
 
 Outlook