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

patito/_pydantic/column_info.py

@@ -1,27 +1,90 @@
 from __future__ import annotations
 
+import io
 import json
-from typing import (
-    Any,
-    Dict,
-    Optional,
-    Sequence,
-    Type,
-    TypeVar,
-    Union,
-)
+from typing import Annotated, Optional, Union
 
 import polars as pl
+from polars.datatypes import *  # noqa: F403 # type: ignore
 from polars.datatypes import DataType, DataTypeClass
-from pydantic import BaseModel, field_serializer
+from polars.exceptions import ComputeError
+from pydantic import BaseModel, BeforeValidator, field_serializer
 
-from patito._pydantic.dtypes import parse_composite_dtype
+
+def dtype_deserializer(dtype: str | DataTypeClass | DataType | None):
+    """Deserialize a dtype from json."""
+    if isinstance(dtype, DataTypeClass) or isinstance(dtype, DataType):
+        return dtype
+    else:
+        if dtype == "null" or dtype is None:
+            return None
+        else:
+            return eval(dtype)
+
+
+def expr_deserializer(
+    expr: str | pl.Expr | list[pl.Expr] | None,
+) -> pl.Expr | list[pl.Expr] | None:
+    """Deserialize a polars expression or list thereof from json.
+
+    This is applied both during deserialization and validation.
+    """
+    if expr is None:
+        return None
+    elif isinstance(expr, pl.Expr):
+        return expr
+    elif isinstance(expr, list):
+        return expr
+    elif isinstance(expr, str):
+        if expr == "null":
+            return None
+        # can be either a list of expr or expr
+        elif expr[0] == "[":
+            return [
+                pl.Expr.deserialize(io.StringIO(e), format="json")
+                for e in json.loads(expr)
+            ]
+        else:
+            return pl.Expr.deserialize(io.StringIO(expr), format="json")
+    else:
+        raise ValueError(f"{expr} can not be deserialized.")
+
+
+def expr_or_col_name_deserializer(expr: str | pl.Expr | None) -> pl.Expr | str | None:
+    """Deserialize a polars expression or column name from json.
+
+    This is applied both during deserialization and validation.
+    """
+    if expr is None:
+        return None
+    elif isinstance(expr, pl.Expr):
+        return expr
+    elif isinstance(expr, list):
+        return expr
+    elif isinstance(expr, str):
+        # Default behaviour
+        if expr == "null":
+            return None
+        else:
+            try:
+                return pl.Expr.deserialize(io.StringIO(expr), format="json")
+            except ComputeError:
+                try:
+                    # Column name is being deserialized
+                    return json.loads(expr)
+                except json.JSONDecodeError:
+                    # Column name has been passed literally
+                    # to ColumnInfo(derived_from="foo")
+                    return expr
+    else:
+        raise ValueError(f"{expr} can not be deserialized.")
 
 
 class ColumnInfo(BaseModel, arbitrary_types_allowed=True):
     """patito-side model for storing column metadata.
 
     Args:
+        allow_missing (bool): Column may be missing.
         constraints (Union[polars.Expression, List[polars.Expression]): A single
             constraint or list of constraints, expressed as a polars expression objects.
             All rows must satisfy the given constraint. You can refer to the given column
@@ -34,10 +97,20 @@ class ColumnInfo(BaseModel, arbitrary_types_allowed=True):
 
     """
 
-    dtype: Optional[Union[DataTypeClass, DataType]] = None
-    constraints: Optional[Union[pl.Expr, Sequence[pl.Expr]]] = None
-    derived_from: Optional[Union[str, pl.Expr]] = None
-    unique: Optional[bool] = None
+    allow_missing: Optional[bool] = None  # noqa: UP007
+    dtype: Annotated[
+        Optional[Union[DataTypeClass, DataType]],  # noqa: UP007
+        BeforeValidator(dtype_deserializer),
+    ] = None
+    constraints: Annotated[
+        Optional[Union[pl.Expr, list[pl.Expr]]],  # noqa: UP007
+        BeforeValidator(expr_deserializer),
+    ] = None
+    derived_from: Annotated[
+        Optional[Union[str, pl.Expr]],  # noqa: UP007
+        BeforeValidator(expr_or_col_name_deserializer),
+    ] = None
+    unique: Optional[bool] = None  # noqa : UP007
 
     def __repr__(self) -> str:
         """Print only Field attributes whose values are not default (mainly None)."""
@@ -56,39 +129,21 @@ class ColumnInfo(BaseModel, arbitrary_types_allowed=True):
         return f"ColumnInfo({string})"
 
     @field_serializer("constraints", "derived_from")
-    def serialize_exprs(self, exprs: str | pl.Expr | Sequence[pl.Expr] | None) -> Any:
-        if exprs is None:
-            return None
-        elif isinstance(exprs, str):
-            return exprs
-        elif isinstance(exprs, pl.Expr):
-            return self._serialize_expr(exprs)
-        elif isinstance(exprs, Sequence):
-            return [self._serialize_expr(c) for c in exprs]
+    def expr_serializer(self, expr: None | pl.Expr | list[pl.Expr]):
+        """Converts polars expr to json."""
+        if expr is None:
+            return "null"
+        elif isinstance(expr, str):
+            return json.dumps(expr)
+        elif isinstance(expr, list):
+            return json.dumps([e.meta.serialize(format="json") for e in expr])
         else:
-            raise ValueError(f"Invalid type for exprs: {type(exprs)}")
-
-    def _serialize_expr(self, expr: pl.Expr) -> Dict:
-        if isinstance(expr, pl.Expr):
-            return json.loads(
-                expr.meta.serialize(format="json")
-            )  # can we access the dictionary directly?
-        else:
-            raise ValueError(f"Invalid type for expr: {type(expr)}")
+            return expr.meta.serialize(format="json")
 
     @field_serializer("dtype")
-    def serialize_dtype(self, dtype: DataTypeClass | DataType | None) -> Any:
-        """Serialize a polars dtype.
-
-        References:
-            [1] https://stackoverflow.com/questions/76572310/how-to-serialize-deserialize-polars-datatypes
-        """
+    def dtype_serializer(self, dtype: DataTypeClass | DataType | None) -> str:
+        """Converts polars dtype to json."""
         if dtype is None:
-            return None
-        elif isinstance(dtype, DataTypeClass) or isinstance(dtype, DataType):
-            return parse_composite_dtype(dtype)
+            return "null"
         else:
-            raise ValueError(f"Invalid type for dtype: {type(dtype)}")
-
-
-CI = TypeVar("CI", bound=Type[ColumnInfo])
+            return str(dtype)

patito/_pydantic/dtypes/dtypes.py

@@ -1,21 +1,22 @@
 from __future__ import annotations
 
+from collections.abc import Mapping
 from functools import cache, reduce
 from operator import and_
-from typing import TYPE_CHECKING, Any, Dict, FrozenSet, Mapping, Optional, Type
+from typing import TYPE_CHECKING, Any
 
 import polars as pl
 from polars.datatypes import DataType, DataTypeClass
 from polars.datatypes.group import DataTypeGroup
 from pydantic import TypeAdapter
 
+from patito._pydantic.column_info import ColumnInfo
 from patito._pydantic.dtypes.utils import (
     PT_BASE_SUPPORTED_DTYPES,
     PydanticBaseType,
     _pyd_type_to_default_dtype,
     _pyd_type_to_valid_dtypes,
     _without_optional,
-    dtype_from_string,
 )
 from patito._pydantic.repr import display_as_type
 
@@ -25,8 +26,8 @@ if TYPE_CHECKING:
 
 @cache
 def valid_dtypes_for_model(
-    cls: Type[ModelType],
-) -> Mapping[str, FrozenSet[DataTypeClass]]:
+    cls: type[ModelType],
+) -> Mapping[str, frozenset[DataTypeClass]]:
     return {
         column: (
             DtypeResolver(cls.model_fields[column].annotation).valid_polars_dtypes()
@@ -39,7 +40,7 @@ def valid_dtypes_for_model(
 
 @cache
 def default_dtypes_for_model(
-    cls: Type[ModelType],
+    cls: type[ModelType],
 ) -> dict[str, DataType]:
     default_dtypes: dict[str, DataType] = {}
     for column in cls.columns:
@@ -57,7 +58,7 @@ def default_dtypes_for_model(
 def validate_polars_dtype(
     annotation: type[Any] | None,
     dtype: DataType | DataTypeClass | None,
-    column: Optional[str] = None,
+    column: str | None = None,
 ) -> None:
     """Check that the polars dtype is valid for the given annotation. Raises ValueError if not.
 
@@ -84,7 +85,7 @@ def validate_polars_dtype(
 
 
 def validate_annotation(
-    annotation: type[Any] | Any | None, column: Optional[str] = None
+    annotation: type[Any] | Any | None, column: str | None = None
 ) -> None:
     """Check that the provided annotation has polars/patito support (we can resolve it to a default dtype). Raises ValueError if not.
 
@@ -129,7 +130,7 @@ class DtypeResolver:
 
     def _valid_polars_dtypes_for_schema(
         self,
-        schema: Dict,
+        schema: dict,
     ) -> DataTypeGroup:
         valid_type_sets = []
         if "anyOf" in schema:
@@ -146,7 +147,7 @@ class DtypeResolver:
 
     def _pydantic_subschema_to_valid_polars_types(
         self,
-        props: Dict,
+        props: dict,
     ) -> DataTypeGroup:
         if "type" not in props:
             if "enum" in props:
@@ -189,7 +190,7 @@ class DtypeResolver:
             PydanticBaseType(pyd_type), props.get("format"), props.get("enum")
         )
 
-    def _default_polars_dtype_for_schema(self, schema: Dict) -> DataType | None:
+    def _default_polars_dtype_for_schema(self, schema: dict) -> DataType | None:
         if "anyOf" in schema:
             if len(schema["anyOf"]) == 2:  # look for optionals first
                 schema = _without_optional(schema)
@@ -205,12 +206,12 @@ class DtypeResolver:
 
     def _pydantic_subschema_to_default_dtype(
         self,
-        props: Dict,
+        props: dict,
     ) -> DataType | None:
         if "column_info" in props:  # user has specified in patito model
-            if props["column_info"]["dtype"] is not None:
-                dtype = dtype_from_string(props["column_info"]["dtype"])
-                dtype = dtype() if isinstance(dtype, DataTypeClass) else dtype
+            ci = ColumnInfo.model_validate_json(props["column_info"])
+            if ci.dtype is not None:
+                dtype = ci.dtype() if isinstance(ci.dtype, DataTypeClass) else ci.dtype
                 return dtype
         if "type" not in props:
             if "enum" in props:

patito/_pydantic/dtypes/utils.py

@@ -1,15 +1,11 @@
 from __future__ import annotations
 
 import sys
+from collections.abc import Sequence
 from enum import Enum
 from typing import (
     Any,
-    Dict,
-    List,
-    Optional,
-    Sequence,
     Union,
-    cast,
     get_args,
     get_origin,
 )
@@ -23,9 +19,6 @@ from polars.datatypes.group import (
     INTEGER_DTYPES,
     DataTypeGroup,
 )
-from polars.polars import (
-    dtype_str_repr,  # TODO: this is a rust function, can we implement our own string parser for Time/Duration/Datetime?
-)
 
 PYTHON_TO_PYDANTIC_TYPES = {
     str: "string",
@@ -90,32 +83,42 @@ def is_optional(type_annotation: type[Any] | Any | None) -> bool:
     )
 
 
+def unwrap_optional(type_annotation: type[Any] | Any) -> type:
+    """Return the inner, wrapped type of an Optional.
+
+    Is a no-op for non-Optional types.
+
+    Args:
+        type_annotation: The type annotation to be dewrapped.
+
+    Returns:
+        The input type, but with the outermost Optional removed.
+
+    """
+    return (
+        next(  # pragma: no cover
+            valid_type
+            for valid_type in get_args(type_annotation)
+            if valid_type is not type(None)  # noqa: E721
+        )
+        if is_optional(type_annotation)
+        else type_annotation
+    )
+
+
 def parse_composite_dtype(dtype: DataTypeClass | DataType) -> str:
     """For serialization, converts polars dtype to string representation."""
-    if dtype.is_nested():
-        if dtype == pl.Struct or isinstance(dtype, pl.Struct):
-            raise NotImplementedError("Structs not yet supported by patito")
-        if not isinstance(dtype, pl.List) or isinstance(dtype, pl.Array):
-            raise NotImplementedError(
-                f"Unsupported nested dtype: {dtype} of type {type(dtype)}"
-            )
-        if dtype.inner is None:
-            return convert.DataTypeMappings.DTYPE_TO_FFINAME[dtype.base_type()]
-        return f"{convert.DataTypeMappings.DTYPE_TO_FFINAME[dtype.base_type()]}[{parse_composite_dtype(dtype.inner)}]"
-    elif dtype.is_temporal():
-        return cast(str, dtype_str_repr(dtype))
-    else:
-        return convert.DataTypeMappings.DTYPE_TO_FFINAME[dtype]
+    return str(dtype)
 
 
-def dtype_from_string(v: str) -> Optional[Union[DataTypeClass, DataType]]:
+def dtype_from_string(v: str) -> DataTypeClass | DataType | None:
     """For deserialization."""
     # TODO test all dtypes
     return convert.dtype_short_repr_to_dtype(v)
 
 
 def _pyd_type_to_valid_dtypes(
-    pyd_type: PydanticBaseType, string_format: Optional[str], enum: List[str] | None
+    pyd_type: PydanticBaseType, string_format: str | None, enum: list[str] | None
 ) -> DataTypeGroup:
     if enum is not None:
         _validate_enum_values(pyd_type, enum)
@@ -142,7 +145,7 @@ def _pyd_type_to_valid_dtypes(
 
 
 def _pyd_type_to_default_dtype(
-    pyd_type: PydanticBaseType, string_format: Optional[str], enum: List[str] | None
+    pyd_type: PydanticBaseType, string_format: str | None, enum: list[str] | None
 ) -> DataTypeClass | DataType:
     if enum is not None:
         _validate_enum_values(pyd_type, enum)
@@ -208,7 +211,7 @@ def _pyd_string_format_to_default_dtype(
         raise NotImplementedError
 
 
-def _without_optional(schema: Dict) -> Dict:
+def _without_optional(schema: dict) -> dict:
     if "anyOf" in schema:
         for sub_props in schema["anyOf"]:
             if "type" in sub_props and sub_props["type"] == "null":

patito/_pydantic/repr.py

@@ -1,26 +1,23 @@
 import sys
 import types
 import typing
+from collections.abc import Generator, Iterable, Sequence
 from typing import (
     Any,
     Callable,
-    Generator,
-    Iterable,
     Literal,
     Optional,
-    Sequence,
-    Tuple,
-    Type,
     Union,
     get_args,
     get_origin,
 )
+from typing import GenericAlias as TypingGenericAlias  # type: ignore
 
 if typing.TYPE_CHECKING:
-    Loc = Tuple[Union[int, str], ...]
-    ReprArgs = Sequence[Tuple[Optional[str], Any]]
+    Loc = tuple[Union[int, str], ...]
+    ReprArgs = Sequence[tuple[Optional[str], Any]]
     RichReprResult = Iterable[
-        Union[Any, Tuple[Any], Tuple[str, Any], Tuple[str, Any, Any]]
+        Union[Any, tuple[Any], tuple[str, Any], tuple[str, Any, Any]]
     ]
 
 try:
@@ -30,15 +27,10 @@ except ImportError:
 
 typing_base = _TypingBase
 
-if sys.version_info < (3, 9):
-    # python < 3.9 does not have GenericAlias (list[int], tuple[str, ...] and so on)
-    TypingGenericAlias = ()
-else:
-    from typing import GenericAlias as TypingGenericAlias  # type: ignore
 
 if sys.version_info < (3, 10):
 
-    def origin_is_union(tp: Optional[Type[Any]]) -> bool:
+    def origin_is_union(tp: Optional[type[Any]]) -> bool:
         return tp is typing.Union
 
     WithArgsTypes = (TypingGenericAlias,)
@@ -58,7 +50,7 @@ class Representation:
     of objects.
     """
 
-    __slots__: Tuple[str, ...] = tuple()
+    __slots__: tuple[str, ...] = tuple()
 
     def __repr_args__(self) -> "ReprArgs":
         """Returns the attributes to show in __str__, __repr__, and __pretty__ this is generally overridden.

patito/_pydantic/schema.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+from collections.abc import Mapping
 from functools import cache
-from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Type, cast, get_args
+from typing import TYPE_CHECKING, Any, get_args
 
 from pydantic.fields import FieldInfo
 
@@ -13,7 +14,7 @@ if TYPE_CHECKING:
 
 
 @cache
-def schema_for_model(cls: Type[ModelType]) -> Dict[str, Dict[str, Any]]:
+def schema_for_model(cls: type[ModelType]) -> dict[str, dict[str, Any]]:
     """Return schema properties where definition references have been resolved.
 
     Returns:
@@ -46,27 +47,27 @@ def schema_for_model(cls: Type[ModelType]) -> Dict[str, Dict[str, Any]]:
 
 
 @cache
-def column_infos_for_model(cls: Type[ModelType]) -> Mapping[str, ColumnInfo]:
+def column_infos_for_model(cls: type[ModelType]) -> Mapping[str, ColumnInfo]:
     fields = cls.model_fields
 
     def get_column_info(field: FieldInfo) -> ColumnInfo:
         if field.json_schema_extra is None:
-            return cast(ColumnInfo, cls.column_info_class())
+            return ColumnInfo()
         elif callable(field.json_schema_extra):
             raise NotImplementedError(
                 "Callable json_schema_extra not supported by patito."
             )
-        return cast(ColumnInfo, field.json_schema_extra["column_info"])
+        return ColumnInfo.model_validate_json(field.json_schema_extra["column_info"])
 
     return {k: get_column_info(v) for k, v in fields.items()}
 
 
 def _append_field_info_to_props(
-    field_info: Dict[str, Any],
+    field_info: dict[str, Any],
     field_name: str,
-    model_schema: Dict[str, Any],
-    required: Optional[bool] = None,
-) -> Dict[str, Any]:
+    model_schema: dict[str, Any],
+    required: bool | None = None,
+) -> dict[str, Any]:
     if "$ref" in field_info:  # TODO onto runtime append
         definition = model_schema["$defs"][field_info["$ref"]]
         if "enum" in definition and "type" not in definition:

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"