patito 0.5.1__py3-none-any.whl → 0.6.2__py3-none-any.whl

patito/__init__.py

@@ -1,47 +1,28 @@
 """Patito, a data-modelling library built on top of polars and pydantic."""
+
 from polars import Expr, Series, col
 
-from patito import exceptions, sql
-from patito.exceptions import ValidationError
+from patito import exceptions
+from patito.exceptions import DataFrameValidationError
 from patito.polars import DataFrame, LazyFrame
 from patito.pydantic import Field, Model
 
 _CACHING_AVAILABLE = False
-_DUCKDB_AVAILABLE = False
 field = col("_")
 __all__ = [
     "DataFrame",
+    "DataFrameValidationError",
     "Expr",
     "Field",
     "LazyFrame",
     "Model",
     "Series",
-    "ValidationError",
     "_CACHING_AVAILABLE",
-    "_DUCKDB_AVAILABLE",
     "col",
     "exceptions",
     "field",
-    "sql",
 ]
 
-try:
-    from patito import duckdb
-
-    _DUCKDB_AVAILABLE = True
-    __all__ += ["duckdb"]
-except ImportError:  # pragma: no cover
-    pass
-
-try:
-    from patito.database import Database
-
-    _CACHING_AVAILABLE = True
-    __all__ += ["Database"]
-except ImportError:
-    pass
-
-
 try:
     from importlib.metadata import PackageNotFoundError, version
 except ImportError:  # pragma: no cover

patito/_docs.py

@@ -1,2 +1,3 @@
 """Ugly workaround for Sphinx + autodoc + ModelMetaclass + classproperty."""
+
 from patito.pydantic import ModelMetaclass as Model  # noqa: F401, pragma: no cover

patito/_pydantic/column_info.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import json
+from typing import (
+    Any,
+    Dict,
+    Optional,
+    Sequence,
+    Type,
+    TypeVar,
+    Union,
+)
+
+import polars as pl
+from polars.datatypes import DataType, DataTypeClass
+from pydantic import BaseModel, field_serializer
+
+from patito._pydantic.dtypes import parse_composite_dtype
+
+
+class ColumnInfo(BaseModel, arbitrary_types_allowed=True):
+    """patito-side model for storing column metadata.
+
+    Args:
+        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
+            with ``pt.field``, which will automatically be replaced with
+            ``polars.col(<field_name>)`` before evaluation.
+        derived_from (Union[str, polars.Expr]): used to mark fields that are meant to be derived from other fields. Users can specify a polars expression that will be called to derive the column value when `pt.DataFrame.derive` is called.
+        dtype (polars.datatype.DataType): The given dataframe column must have the given
+            polars dtype, for instance ``polars.UInt64`` or ``pl.Float32``.
+        unique (bool): All row values must be unique.
+
+    """
+
+    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
+
+    def __repr__(self) -> str:
+        """Print only Field attributes whose values are not default (mainly None)."""
+        not_default_field = {
+            field: getattr(self, field)
+            for field in self.model_fields
+            if getattr(self, field) is not self.model_fields[field].default
+        }
+
+        string = ""
+        for field, value in not_default_field.items():
+            string += f"{field}={value}, "
+        if string:
+            # remove trailing comma and space
+            string = string[:-2]
+        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]
+        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(None)
+            )  # can we access the dictionary directly?
+        else:
+            raise ValueError(f"Invalid type for expr: {type(expr)}")
+
+    @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
+        """
+        if dtype is None:
+            return None
+        elif isinstance(dtype, DataTypeClass) or isinstance(dtype, DataType):
+            return parse_composite_dtype(dtype)
+        else:
+            raise ValueError(f"Invalid type for dtype: {type(dtype)}")
+
+
+CI = TypeVar("CI", bound=Type[ColumnInfo])

patito/_pydantic/dtypes/__init__.py

@@ -0,0 +1,25 @@
+from patito._pydantic.dtypes.dtypes import (
+    DtypeResolver,
+    default_dtypes_for_model,
+    valid_dtypes_for_model,
+    validate_annotation,
+    validate_polars_dtype,
+)
+from patito._pydantic.dtypes.utils import (
+    PYTHON_TO_PYDANTIC_TYPES,
+    dtype_from_string,
+    is_optional,
+    parse_composite_dtype,
+)
+
+__all__ = [
+    "DtypeResolver",
+    "validate_annotation",
+    "validate_polars_dtype",
+    "parse_composite_dtype",
+    "dtype_from_string",
+    "valid_dtypes_for_model",
+    "default_dtypes_for_model",
+    "PYTHON_TO_PYDANTIC_TYPES",
+    "is_optional",
+]

patito/_pydantic/dtypes/dtypes.py

@@ -0,0 +1,249 @@
+from __future__ import annotations
+
+from functools import cache, reduce
+from operator import and_
+from typing import TYPE_CHECKING, Any, Dict, FrozenSet, Mapping, Optional, Type
+
+import polars as pl
+from polars.datatypes import DataType, DataTypeClass, DataTypeGroup
+from pydantic import TypeAdapter
+
+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
+
+if TYPE_CHECKING:
+    from patito.pydantic import ModelType
+
+
+@cache
+def valid_dtypes_for_model(
+    cls: Type[ModelType],
+) -> Mapping[str, FrozenSet[DataTypeClass]]:
+    return {
+        column: (
+            DtypeResolver(cls.model_fields[column].annotation).valid_polars_dtypes()
+            if cls.column_infos[column].dtype is None
+            else DataTypeGroup([cls.dtypes[column]], match_base_type=False)
+        )
+        for column in cls.columns
+    }
+
+
+@cache
+def default_dtypes_for_model(
+    cls: Type[ModelType],
+) -> dict[str, DataType]:
+    default_dtypes: dict[str, DataType] = {}
+    for column in cls.columns:
+        dtype = (
+            cls.column_infos[column].dtype
+            or DtypeResolver(cls.model_fields[column].annotation).default_polars_dtype()
+        )
+        if dtype is None:
+            raise ValueError(f"Unable to find a default dtype for column `{column}`")
+
+        default_dtypes[column] = dtype if isinstance(dtype, DataType) else dtype()
+    return default_dtypes
+
+
+def validate_polars_dtype(
+    annotation: type[Any] | None,
+    dtype: DataType | DataTypeClass | None,
+    column: Optional[str] = None,
+) -> None:
+    """Check that the polars dtype is valid for the given annotation. Raises ValueError if not.
+
+    Args:
+        annotation (type[Any] | None): python type annotation
+        dtype (DataType | DataTypeClass | None): polars dtype
+        column (Optional[str], optional): column name. Defaults to None.
+
+    """
+    if (
+        dtype is None or annotation is None
+    ):  # no potential conflict between type annotation and chosen polars type
+        return
+    valid_dtypes = DtypeResolver(annotation).valid_polars_dtypes()
+    if dtype not in valid_dtypes:
+        if column:
+            column_msg = f" for column `{column}`"
+        else:
+            column_msg = ""
+        raise ValueError(
+            f"Invalid dtype {dtype}{column_msg}. Allowable polars dtypes for {display_as_type(annotation)} are: {', '.join([str(x) for x in valid_dtypes])}."
+        )
+    return
+
+
+def validate_annotation(
+    annotation: type[Any] | Any | None, column: Optional[str] = None
+) -> None:
+    """Check that the provided annotation has polars/patito support (we can resolve it to a default dtype). Raises ValueError if not.
+
+    Args:
+        annotation (type[Any] | None): python type annotation
+        column (Optional[str], optional): column name. Defaults to None.
+
+    """
+    default_dtype = DtypeResolver(annotation).default_polars_dtype()
+    if default_dtype is None:
+        valid_polars_dtypes = DtypeResolver(annotation).valid_polars_dtypes()
+        if column:
+            column_msg = f" for column `{column}`"
+        else:
+            column_msg = ""
+        if len(valid_polars_dtypes) == 0:
+            raise ValueError(
+                f"Annotation {display_as_type(annotation)}{column_msg} is not compatible with any polars dtypes."
+            )
+        else:
+            raise ValueError(
+                f"Unable to determine default dtype for annotation {display_as_type(annotation)}{column_msg}. Please provide a valid default polars dtype via the `dtype` argument to `Field`. Valid dtypes are: {', '.join([str(x) for x in valid_polars_dtypes])}."
+            )
+    return
+
+
+class DtypeResolver:
+    def __init__(self, annotation: Any | None):
+        self.annotation = annotation
+        self.schema = TypeAdapter(annotation).json_schema()
+        self.defs = self.schema.get("$defs", {})
+
+    def valid_polars_dtypes(self) -> DataTypeGroup:
+        if self.annotation == Any:
+            return PT_BASE_SUPPORTED_DTYPES
+        return self._valid_polars_dtypes_for_schema(self.schema)
+
+    def default_polars_dtype(self) -> DataType | None:
+        if self.annotation == Any:
+            return pl.String()
+        return self._default_polars_dtype_for_schema(self.schema)
+
+    def _valid_polars_dtypes_for_schema(
+        self,
+        schema: Dict,
+    ) -> DataTypeGroup:
+        valid_type_sets = []
+        if "anyOf" in schema:
+            schema = _without_optional(schema)
+            for sub_props in schema["anyOf"]:
+                valid_type_sets.append(
+                    self._pydantic_subschema_to_valid_polars_types(sub_props)
+                )
+        else:
+            valid_type_sets.append(
+                self._pydantic_subschema_to_valid_polars_types(schema)
+            )
+        return reduce(and_, valid_type_sets) if valid_type_sets else DataTypeGroup([])
+
+    def _pydantic_subschema_to_valid_polars_types(
+        self,
+        props: Dict,
+    ) -> DataTypeGroup:
+        if "type" not in props:
+            if "enum" in props:
+                raise TypeError("Mixed type enums not supported by patito.")
+            elif "const" in props:
+                return DtypeResolver(type(props["const"])).valid_polars_dtypes()
+            elif "$ref" in props:
+                return self._pydantic_subschema_to_valid_polars_types(
+                    self.defs[props["$ref"].split("/")[-1]]
+                )
+            return DataTypeGroup([])
+        pyd_type = props.get("type")
+        if pyd_type == "array":
+            if "items" not in props:
+                return DataTypeGroup([])
+            array_props = props["items"]
+            item_dtypes = self._valid_polars_dtypes_for_schema(array_props)
+            # TODO support pl.Array?
+            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"]
+            return DataTypeGroup(
+                [
+                    pl.Struct(
+                        [
+                            pl.Field(
+                                name, self._default_polars_dtype_for_schema(sub_props)
+                            )
+                            for name, sub_props in object_props.items()
+                        ]
+                    )
+                ],
+                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:
+        if "anyOf" in schema:
+            if len(schema["anyOf"]) == 2:  # look for optionals first
+                schema = _without_optional(schema)
+            if len(schema["anyOf"]) == 1:
+                if "column_info" in schema:
+                    schema["anyOf"][0]["column_info"] = schema[
+                        "column_info"
+                    ]  # push column info through optional
+                schema = schema["anyOf"][0]
+            else:
+                return None
+        return self._pydantic_subschema_to_default_dtype(schema)
+
+    def _pydantic_subschema_to_default_dtype(
+        self,
+        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
+                return dtype
+        if "type" not in props:
+            if "enum" in props:
+                raise TypeError("Mixed type enums not supported by patito.")
+            elif "const" in props:
+                return DtypeResolver(type(props["const"])).default_polars_dtype()
+            elif "$ref" in props:
+                return self._pydantic_subschema_to_default_dtype(
+                    self.defs[props["$ref"].split("/")[-1]]
+                )
+            return None
+        pyd_type = props.get("type")
+        if pyd_type == "array":
+            if "items" not in props:
+                raise NotImplementedError(
+                    "Unexpected error processing pydantic schema. Please file an issue."
+                )
+            array_props = props["items"]
+            inner_default_type = self._default_polars_dtype_for_schema(array_props)
+            if inner_default_type is None:
+                return None
+            return pl.List(inner_default_type)
+        elif pyd_type == "object":
+            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()
+                ]
+            )
+        return _pyd_type_to_default_dtype(
+            PydanticBaseType(pyd_type), props.get("format"), props.get("enum")
+        )

patito/_pydantic/dtypes/utils.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+import sys
+from enum import Enum
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Union,
+    cast,
+    get_args,
+    get_origin,
+)
+
+import polars as pl
+from polars.datatypes import DataType, DataTypeClass, DataTypeGroup, convert
+from polars.datatypes.constants import (
+    DATETIME_DTYPES,
+    DURATION_DTYPES,
+    FLOAT_DTYPES,
+    INTEGER_DTYPES,
+)
+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",
+    int: "integer",
+    float: "number",
+    bool: "boolean",
+    type(None): "null",
+}
+
+BOOLEAN_DTYPES = DataTypeGroup([pl.Boolean])
+STRING_DTYPES = DataTypeGroup([pl.String])
+DATE_DTYPES = DataTypeGroup([pl.Date])
+TIME_DTYPES = DataTypeGroup([pl.Time])
+
+PT_BASE_SUPPORTED_DTYPES = DataTypeGroup(
+    INTEGER_DTYPES
+    | FLOAT_DTYPES
+    | BOOLEAN_DTYPES
+    | STRING_DTYPES
+    | DATE_DTYPES
+    | DATETIME_DTYPES
+    | DURATION_DTYPES
+    | TIME_DTYPES
+)
+
+if sys.version_info >= (3, 10):  # pragma: no cover
+    from types import UnionType  # pyright: ignore
+
+    UNION_TYPES = (Union, UnionType)
+else:
+    UNION_TYPES = (Union,)  # pragma: no cover
+
+
+class PydanticBaseType(Enum):
+    STRING = "string"
+    INTEGER = "integer"
+    NUMBER = "number"
+    BOOLEAN = "boolean"
+    NULL = "null"
+    OBJECT = "object"
+
+
+class PydanticStringFormat(Enum):
+    DATE = "date"
+    DATE_TIME = "date-time"
+    DURATION = "duration"
+    TIME = "time"
+
+
+def is_optional(type_annotation: type[Any] | Any | None) -> bool:
+    """Return True if the given type annotation is an Optional annotation.
+
+    Args:
+        type_annotation: The type annotation to be checked.
+
+    Returns:
+        True if the outermost type is Optional.
+
+    """
+    return (get_origin(type_annotation) in UNION_TYPES) and (
+        type(None) in get_args(type_annotation)
+    )
+
+
+def parse_composite_dtype(dtype: DataTypeClass | DataType) -> str:
+    """For serialization, converts polars dtype to string representation."""
+    if dtype in pl.NESTED_DTYPES:
+        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 in pl.TEMPORAL_DTYPES:
+        return cast(str, dtype_str_repr(dtype))
+    else:
+        return convert.DataTypeMappings.DTYPE_TO_FFINAME[dtype]
+
+
+def dtype_from_string(v: str) -> Optional[Union[DataTypeClass, DataType]]:
+    """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
+) -> 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)
+    elif pyd_type.value == "number":
+        return (
+            FLOAT_DTYPES
+            if isinstance(FLOAT_DTYPES, DataTypeGroup)
+            else DataTypeGroup(FLOAT_DTYPES)
+        )
+    elif pyd_type.value == "boolean":
+        return BOOLEAN_DTYPES
+    elif pyd_type.value == "string":
+        _string_format = (
+            PydanticStringFormat(string_format) if string_format is not None else None
+        )
+        return _pyd_string_format_to_valid_dtypes(_string_format)
+    elif pyd_type.value == "null":
+        return DataTypeGroup([pl.Null])
+    else:
+        return DataTypeGroup([])
+
+
+def _pyd_type_to_default_dtype(
+    pyd_type: PydanticBaseType, string_format: Optional[str], enum: List[str] | None
+) -> DataTypeClass | DataType:
+    if enum is not None:
+        _validate_enum_values(pyd_type, enum)
+        return pl.Enum(enum)
+    elif pyd_type.value == "integer":
+        return pl.Int64()
+    elif pyd_type.value == "number":
+        return pl.Float64()
+    elif pyd_type.value == "boolean":
+        return pl.Boolean()
+    elif pyd_type.value == "string":
+        _string_format = (
+            PydanticStringFormat(string_format) if string_format is not None else None
+        )
+        return _pyd_string_format_to_default_dtype(_string_format)
+    elif pyd_type.value == "null":
+        return pl.Null()
+    elif pyd_type.value == "object":
+        raise ValueError("pydantic object types not currently supported by patito")
+    else:
+        raise NotImplementedError
+
+
+def _pyd_string_format_to_valid_dtypes(
+    string_format: PydanticStringFormat | None,
+) -> DataTypeGroup:
+    if string_format is None:
+        return STRING_DTYPES
+    elif string_format.value == "date":
+        return DATE_DTYPES
+    elif string_format.value == "date-time":
+        return (
+            DATETIME_DTYPES
+            if isinstance(DATE_DTYPES, DataTypeGroup)
+            else DataTypeGroup(DATE_DTYPES)
+        )
+    elif string_format.value == "duration":
+        return (
+            DURATION_DTYPES
+            if isinstance(DURATION_DTYPES, DataTypeGroup)
+            else DataTypeGroup(DURATION_DTYPES)
+        )
+    elif string_format.value == "time":
+        return TIME_DTYPES
+    else:
+        raise NotImplementedError
+
+
+def _pyd_string_format_to_default_dtype(
+    string_format: PydanticStringFormat | None,
+) -> DataTypeClass | DataType:
+    if string_format is None:
+        return pl.String()
+    elif string_format.value == "date":
+        return pl.Date()
+    elif string_format.value == "date-time":
+        return pl.Datetime()
+    elif string_format.value == "duration":
+        return pl.Duration()
+    elif string_format.value == "time":
+        return pl.Time()
+    else:
+        raise NotImplementedError
+
+
+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":
+                schema["anyOf"].remove(sub_props)
+    return schema
+
+
+def _validate_enum_values(pyd_type: PydanticBaseType, enum: Sequence[Any]) -> None:
+    enum_types = set(type(value) for value in enum)
+    if len(enum_types) > 1:
+        raise TypeError(
+            f"All enumerated values of enums used to annotate Patito model fields must have the same type. Encountered types: {sorted(map(lambda t: t.__name__, enum_types))}."
+        )
+    if pyd_type.value != "string":
+        raise TypeError(
+            f"Enums used to annotate Patito model fields must be strings. Encountered type: {enum_types.pop().__name__}."
+        )