patito 0.7.0__py3-none-any.whl → 0.8.2__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,9 +97,19 @@ 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
+    allow_missing: Optional[bool] = None
+    dtype: Annotated[
+        Optional[Union[DataTypeClass, DataType]],
+        BeforeValidator(dtype_deserializer),
+    ] = None
+    constraints: Annotated[
+        Optional[Union[pl.Expr, list[pl.Expr]]],
+        BeforeValidator(expr_deserializer),
+    ] = None
+    derived_from: Annotated[
+        Optional[Union[str, pl.Expr]],
+        BeforeValidator(expr_or_col_name_deserializer),
+    ] = None
     unique: Optional[bool] = None
 
     def __repr__(self) -> str:
@@ -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 operator import or_
+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.
 
@@ -114,7 +115,8 @@ def validate_annotation(
 class DtypeResolver:
     def __init__(self, annotation: Any | None):
         self.annotation = annotation
-        self.schema = TypeAdapter(annotation).json_schema()
+        # mode='serialization' allows nested models with structs, see #86
+        self.schema = TypeAdapter(annotation).json_schema(mode="serialization")
         self.defs = self.schema.get("$defs", {})
 
     def valid_polars_dtypes(self) -> DataTypeGroup:
@@ -129,7 +131,7 @@ class DtypeResolver:
 
     def _valid_polars_dtypes_for_schema(
         self,
-        schema: Dict,
+        schema: dict,
     ) -> DataTypeGroup:
         valid_type_sets = []
         if "anyOf" in schema:
@@ -142,11 +144,11 @@ class DtypeResolver:
             valid_type_sets.append(
                 self._pydantic_subschema_to_valid_polars_types(schema)
             )
-        return reduce(and_, valid_type_sets) if valid_type_sets else DataTypeGroup([])
+        return reduce(or_, valid_type_sets) if valid_type_sets else DataTypeGroup([])
 
     def _pydantic_subschema_to_valid_polars_types(
         self,
-        props: Dict,
+        props: dict,
     ) -> DataTypeGroup:
         if "type" not in props:
             if "enum" in props:
@@ -158,6 +160,7 @@ class DtypeResolver:
                     self.defs[props["$ref"].split("/")[-1]]
                 )
             return DataTypeGroup([])
+
         pyd_type = props.get("type")
         if pyd_type == "array":
             if "items" not in props:
@@ -168,28 +171,27 @@ class DtypeResolver:
             return DataTypeGroup(
                 [pl.List(dtype) for dtype in item_dtypes], match_base_type=False
             )
+
         elif pyd_type == "object":
             if "properties" not in props:
                 return DataTypeGroup([])
             object_props = props["properties"]
+            struct_fields: list[pl.Field] = []
+            for name, sub_props in object_props.items():
+                dtype = self._default_polars_dtype_for_schema(sub_props)
+                assert dtype is not None
+                struct_fields.append(pl.Field(name, dtype))
             return DataTypeGroup(
-                [
-                    pl.Struct(
-                        [
-                            pl.Field(
-                                name, self._default_polars_dtype_for_schema(sub_props)
-                            )
-                            for name, sub_props in object_props.items()
-                        ]
-                    )
-                ],
+                [pl.Struct(struct_fields)],
                 match_base_type=False,
             )  # for structs, return only the default dtype set to avoid combinatoric issues
         return _pyd_type_to_valid_dtypes(
             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[str, Any]
+    ) -> DataType | None:
         if "anyOf" in schema:
             if len(schema["anyOf"]) == 2:  # look for optionals first
                 schema = _without_optional(schema)
@@ -205,13 +207,14 @@ class DtypeResolver:
 
     def _pydantic_subschema_to_default_dtype(
         self,
-        props: Dict,
+        props: dict[str, Any],
     ) -> 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:
                 raise TypeError("Mixed type enums not supported by patito.")
@@ -222,10 +225,12 @@ class DtypeResolver:
                     self.defs[props["$ref"].split("/")[-1]]
                 )
             return None
+
         pyd_type = props.get("type")
         if pyd_type == "numeric":
             pyd_type = "number"
-        if pyd_type == "array":
+
+        elif pyd_type == "array":
             if "items" not in props:
                 raise NotImplementedError(
                     "Unexpected error processing pydantic schema. Please file an issue."
@@ -235,18 +240,21 @@ class DtypeResolver:
             if inner_default_type is None:
                 return None
             return pl.List(inner_default_type)
-        elif pyd_type == "object":
+
+        elif pyd_type == "object":  # these are structs
             if "properties" not in props:
                 raise NotImplementedError(
                     "dictionaries not currently supported by patito"
                 )
-            object_props = props["properties"]
-            return pl.Struct(
-                [
-                    pl.Field(name, self._default_polars_dtype_for_schema(sub_props))
-                    for name, sub_props in object_props.items()
-                ]
-            )
+            object_props: dict[str, dict[str, str]] = props["properties"]
+            struct_fields: list[pl.Field] = []
+
+            for name, sub_props in object_props.items():
+                dtype = self._default_polars_dtype_for_schema(sub_props)
+                assert dtype is not None
+                struct_fields.append(pl.Field(name, dtype))
+            return pl.Struct(struct_fields)
+
         return _pyd_type_to_default_dtype(
             PydanticBaseType(pyd_type), props.get("format"), props.get("enum")
         )

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,38 +83,48 @@ 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)
         return DataTypeGroup([pl.Enum(enum), pl.String], match_base_type=False)
     if pyd_type.value == "integer":
-        return DataTypeGroup(INTEGER_DTYPES | FLOAT_DTYPES)
+        return DataTypeGroup(INTEGER_DTYPES)
     elif pyd_type.value == "number":
         return (
             FLOAT_DTYPES
@@ -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: