dataframely 1.1.0__tar.gz → 1.2.1__tar.gz

{dataframely-1.1.0 → dataframely-1.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataframely
-Version: 1.1.0
+Version: 1.2.1
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12

{dataframely-1.1.0 → dataframely-1.2.1}/dataframely/__init__.py

@@ -34,6 +34,7 @@ from .columns import (
     Int64,
     Integer,
     List,
+    Object,
     String,
     Struct,
     Time,
@@ -89,4 +90,5 @@ __all__ = [
     "String",
     "Struct",
     "List",
+    "Object",
 ]

{dataframely-1.1.0 → dataframely-1.2.1}/dataframely/_base_collection.py

@@ -7,7 +7,7 @@ import typing
 from abc import ABCMeta
 from collections.abc import Iterable
 from dataclasses import dataclass, field
-from typing import Annotated, Any, Self, get_args, get_origin
+from typing import Annotated, Any, Self, cast, get_args, get_origin
 
 import polars as pl
 
@@ -183,59 +183,9 @@ class CollectionMeta(ABCMeta):
         # Get all members via the annotations
         if "__annotations__" in source:
             for attr, kls in source["__annotations__"].items():
-                origin = get_origin(kls)
-
-                # optional annotation
-                collection_member = CollectionMember()
-
-                if origin is Annotated:
-                    annotation_args = get_args(kls)
-                    origin_arg0 = get_origin(annotation_args[0])
-                    if not origin_arg0 or not issubclass(origin_arg0, TypedLazyFrame):
-                        raise AnnotationImplementationError(attr, kls)
-                    if len(annotation_args) > 2:
-                        raise AnnotationImplementationError(attr, kls)
-                    if not isinstance(annotation_args[1], CollectionMember):
-                        raise AnnotationImplementationError(attr, kls)
-
-                    # Continue with wrapped FrameType
-                    collection_member = annotation_args[1]
-                    kls = annotation_args[0]
-                    origin = origin_arg0
-
-                if origin is None:
-                    # `None` annotation is not allowed
-                    raise AnnotationImplementationError(attr, kls)
-                elif origin == typing.Union:
-                    # Happy path: optional member
-                    union_args = get_args(kls)
-                    if len(union_args) != 2:
-                        raise AnnotationImplementationError(attr, kls)
-                    if not any(get_origin(arg) is None for arg in union_args):
-                        raise AnnotationImplementationError(attr, kls)
-
-                    [not_none_arg] = [
-                        arg for arg in union_args if get_origin(arg) is not None
-                    ]
-                    if not issubclass(get_origin(not_none_arg), TypedLazyFrame):
-                        raise AnnotationImplementationError(attr, kls)
-
-                    result.members[attr] = MemberInfo(
-                        schema=get_args(not_none_arg)[0],
-                        is_optional=True,
-                        ignored_in_filters=collection_member.ignored_in_filters,
-                    )
-                elif issubclass(origin, TypedLazyFrame):
-                    # Happy path: required member
-                    result.members[attr] = MemberInfo(
-                        schema=get_args(kls)[0],
-                        is_optional=False,
-                        ignored_in_filters=collection_member.ignored_in_filters,
-                        inline_for_sampling=collection_member.inline_for_sampling,
-                    )
-                else:
-                    # Some other unknown annotation
-                    raise AnnotationImplementationError(attr, kls)
+                result.members[attr] = CollectionMeta._derive_member_info(
+                    attr, kls, CollectionMember()
+                )
 
         # Get all filters by traversing the source
         for attr, value in {
@@ -246,6 +196,55 @@ class CollectionMeta(ABCMeta):
 
         return result
 
+    @staticmethod
+    def _derive_member_info(
+        attr: str, type_annotation: Any, collection_member: CollectionMember
+    ) -> MemberInfo:
+        origin = get_origin(type_annotation)
+
+        if origin is None:
+            # `None` annotation is not allowed
+            raise AnnotationImplementationError(attr, type_annotation)
+        elif origin == Annotated:
+            # Maybe happy path: annotated member, dispatch recursively
+            annotation_args = cast(list[Any], get_args(type_annotation))
+            if len(annotation_args) > 2:
+                raise AnnotationImplementationError(attr, type_annotation)
+            if not isinstance(annotation_args[1], CollectionMember):
+                raise AnnotationImplementationError(attr, type_annotation)
+            return CollectionMeta._derive_member_info(
+                attr, annotation_args[0], annotation_args[1]
+            )
+        elif origin == typing.Union:
+            # Happy path: optional member
+            union_args = get_args(type_annotation)
+            if len(union_args) != 2:
+                raise AnnotationImplementationError(attr, type_annotation)
+            if not any(get_origin(arg) is None for arg in union_args):
+                raise AnnotationImplementationError(attr, type_annotation)
+
+            [not_none_arg] = [arg for arg in union_args if get_origin(arg) is not None]
+            if not issubclass(get_origin(not_none_arg), TypedLazyFrame):
+                raise AnnotationImplementationError(attr, type_annotation)
+
+            return MemberInfo(
+                schema=get_args(not_none_arg)[0],
+                is_optional=True,
+                ignored_in_filters=collection_member.ignored_in_filters,
+                inline_for_sampling=collection_member.inline_for_sampling,
+            )
+        elif issubclass(origin, TypedLazyFrame):
+            # Happy path: required member
+            return MemberInfo(
+                schema=get_args(type_annotation)[0],
+                is_optional=False,
+                ignored_in_filters=collection_member.ignored_in_filters,
+                inline_for_sampling=collection_member.inline_for_sampling,
+            )
+        else:
+            # Some other unknown annotation
+            raise AnnotationImplementationError(attr, type_annotation)
+
 
 class BaseCollection(metaclass=CollectionMeta):
     """Internal utility abstraction to reference collections without introducing

{dataframely-1.1.0 → dataframely-1.2.1}/dataframely/_rule.py

@@ -126,5 +126,7 @@ def _with_group_rules(lf: pl.LazyFrame, rules: dict[str, GroupRule]) -> pl.LazyF
     #    preserves the order of the left data frame.
     result = lf
     for group_columns, frame in group_evaluations.items():
-        result = result.join(frame, on=list(group_columns), how="left")
+        result = result.join(
+            frame, on=list(group_columns), how="left", nulls_equal=True
+        )
     return result

{dataframely-1.1.0 → dataframely-1.2.1}/dataframely/collection.py

@@ -4,7 +4,7 @@
 import sys
 import warnings
 from abc import ABC
-from collections.abc import Mapping, MutableMapping, Sequence
+from collections.abc import Mapping, Sequence
 from pathlib import Path
 from typing import Any, Generic, Self, TypeVar, cast
 
@@ -20,10 +20,10 @@ from .random import Generator
 
 if sys.version_info >= (3, 13):
     SamplingType = TypeVar(
-        "SamplingType", bound=MutableMapping[str, Any], default=dict[str, Any]
+        "SamplingType", bound=Mapping[str, Any], default=Mapping[str, Any]
     )
 else:  # pragma: no cover
-    SamplingType = TypeVar("SamplingType", bound=MutableMapping[str, Any])
+    SamplingType = TypeVar("SamplingType", bound=Mapping[str, Any])
 
 
 class Collection(BaseCollection, ABC, Generic[SamplingType]):

{dataframely-1.1.0 → dataframely-1.2.1}/dataframely/columns/__init__.py

@@ -10,6 +10,7 @@ from .enum import Enum
 from .float import Float, Float32, Float64
 from .integer import Int8, Int16, Int32, Int64, Integer, UInt8, UInt16, UInt32, UInt64
 from .list import List
+from .object import Object
 from .string import String
 from .struct import Struct
 
@@ -31,6 +32,7 @@ __all__ = [
     "Int32",
     "Int64",
     "Integer",
+    "Object",
     "UInt8",
     "UInt16",
     "UInt32",

{dataframely-1.1.0 → dataframely-1.2.1}/dataframely/columns/enum.py

@@ -49,7 +49,7 @@ class Enum(Column):
             alias=alias,
             metadata=metadata,
         )
-        self.categories = categories
+        self.categories = list(categories)
 
     @property
     def dtype(self) -> pl.DataType:

dataframely-1.2.1/dataframely/columns/object.py

@@ -0,0 +1,63 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import polars as pl
+
+from dataframely._compat import pa, sa, sa_TypeEngine
+from dataframely.random import Generator
+
+from ._base import Column
+
+
+class Object(Column):
+    """A Python Object column."""
+
+    def __init__(
+        self,
+        *,
+        nullable: bool = True,
+        primary_key: bool = False,
+        check: Callable[[pl.Expr], pl.Expr] | None = None,
+        alias: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ):
+        """
+        Args:
+            nullable: Whether this column may contain null values.
+            primary_key: Whether this column is part of the primary key of the schema.
+            check: A custom check to run for this column. Must return a non-aggregated
+                boolean expression.
+            alias: An overwrite for this column's name which allows for using a column
+                name that is not a valid Python identifier. Especially note that setting
+                this option does _not_ allow to refer to the column with two different
+                names, the specified alias is the only valid name.
+            metadata: A dictionary of metadata to attach to the column.
+        """
+        super().__init__(
+            nullable=nullable,
+            primary_key=primary_key,
+            check=check,
+            alias=alias,
+            metadata=metadata,
+        )
+
+    @property
+    def dtype(self) -> pl.DataType:
+        return pl.Object()
+
+    def sqlalchemy_dtype(self, dialect: sa.Dialect) -> sa_TypeEngine:
+        raise NotImplementedError("SQL column cannot have 'Object' type.")
+
+    @property
+    def pyarrow_dtype(self) -> pa.DataType:
+        raise NotImplementedError("PyArrow column cannot have 'Object' type.")
+
+    def _sample_unchecked(self, generator: Generator, n: int) -> pl.Series:
+        raise NotImplementedError(
+            "Random data sampling not implemented for 'Object' type."
+        )

{dataframely-1.1.0 → dataframely-1.2.1}/dataframely/functional.py

@@ -1,17 +1,21 @@
 # Copyright (c) QuantCo 2025-2025
 # SPDX-License-Identifier: BSD-3-Clause
 
+from collections.abc import Sequence
 from typing import TypeVar
 
 import polars as pl
 
+from ._base_collection import BaseCollection
 from ._typing import LazyFrame
-from .collection import Collection
 from .schema import Schema
 
 S = TypeVar("S", bound=Schema)
 T = TypeVar("T", bound=Schema)
-C = TypeVar("C", bound=Collection)
+
+# NOTE: Binding to `BaseCollection` is required here as the TypeVar default for the
+#  sampling type otherwise causes issues for Python 3.13.
+C = TypeVar("C", bound=BaseCollection)
 
 # ------------------------------------------------------------------------------------ #
 #                                        FILTER                                        #
@@ -62,7 +66,7 @@ def filter_relationship_one_to_at_least_one(
 # ------------------------------------------------------------------------------------ #
 
 
-def concat_collection_members(collections: list[C], /) -> dict[str, pl.LazyFrame]:
+def concat_collection_members(collections: Sequence[C], /) -> dict[str, pl.LazyFrame]:
     """Concatenate the members of collections with the same type.
 
     Args:

{dataframely-1.1.0 → dataframely-1.2.1}/pyproject.toml

@@ -19,7 +19,7 @@ description = "A declarative, polars-native data frame validation library"
 name = "dataframely"
 readme = "README.md"
 requires-python = ">=3.11"
-version = "1.1.0"
+version = "1.2.1"
 
 [project.urls]
 Repository = "https://github.com/quantco/dataframely"

{dataframely-1.1.0 → dataframely-1.2.1}/tests/collection/test_implementation.py

@@ -81,6 +81,24 @@ def test_annotation_union_conflicting_types_failure() -> None:
         )
 
 
+def test_annotation_annotated_success() -> None:
+    """When we use an Annotated type, it must accept a union type."""
+    create_collection_raw(
+        "test",
+        {
+            "first": Annotated[
+                dy.LazyFrame[MyTestSchema] | None, dy.CollectionMember()
+            ],
+        },
+    )
+    create_collection_raw(
+        "test",
+        {
+            "first": dy.LazyFrame[MyTestSchema] | None,
+        },
+    )
+
+
 def test_annotation_only_none_failure() -> None:
     """Annotations must not just be None."""
     with pytest.raises(AnnotationImplementationError):

{dataframely-1.1.0 → dataframely-1.2.1}/tests/collection/test_sample.py

@@ -54,6 +54,21 @@ class MyInlinedCollection(dy.Collection):
         return sample
 
 
+class MyInlinedCollectionWithOptional(dy.Collection):
+    first: Annotated[
+        dy.LazyFrame[MyFirstSchema] | None,
+        dy.CollectionMember(inline_for_sampling=True),
+    ]
+    second: dy.LazyFrame[MySecondSchema]
+
+    @classmethod
+    def _preprocess_sample(
+        cls, sample: dict[str, Any], index: int, generator: Generator
+    ) -> dict[str, Any]:
+        sample["a"] = index
+        return sample
+
+
 class SmallCollection(dy.Collection):
     first: dy.LazyFrame[MyFirstSchema]
 
@@ -117,13 +132,20 @@ def test_sample_with_overrides() -> None:
     assert collection.second.collect()["c"].to_list() == [3, 4, 6]
 
 
-def test_sample_inline_with_overrides() -> None:
-    collection = MyInlinedCollection.sample(
+@pytest.mark.parametrize(
+    "collection_type", [MyInlinedCollection, MyInlinedCollectionWithOptional]
+)
+def test_sample_inline_with_overrides(
+    collection_type: type[MyInlinedCollection] | type[MyInlinedCollectionWithOptional],
+) -> None:
+    collection = collection_type.sample(
         overrides=[
             {"b": 4, "second": [{"c": 3}, {"c": 4}]},
             {"b": 8, "second": [{"c": 6}]},
         ]
     )
+
+    assert collection.first is not None
     assert collection.first.collect()["a"].to_list() == [0, 1]
     assert collection.first.collect()["b"].to_list() == [4, 8]
 

{dataframely-1.1.0 → dataframely-1.2.1}/tests/column_types/test_enum.py

@@ -52,3 +52,12 @@ def test_valid_cast(
     schema = create_schema("test", {"a": enum})
     df = df_type(data)
     assert schema.is_valid(df, cast=True) == valid
+
+
+@pytest.mark.parametrize("type1", [list, tuple])
+@pytest.mark.parametrize("type2", [list, tuple])
+def test_different_sequences(type1: type, type2: type) -> None:
+    allowed = ["a", "b"]
+    S = create_schema("test", {"x": dy.Enum(type1(allowed))})
+    df = pl.DataFrame({"x": pl.Series(["a", "b"], dtype=pl.Enum(type2(allowed)))})
+    S.validate(df)

dataframely-1.2.1/tests/column_types/test_object.py

@@ -0,0 +1,60 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+
+import polars as pl
+import pytest
+
+import dataframely as dy
+from dataframely.columns._base import Column
+from dataframely.random import Generator
+from dataframely.testing import create_schema
+
+
+class CustomObject:
+    def __init__(self, a: int, b: str) -> None:
+        self.a = a
+        self.b = b
+
+
+def test_simple_object() -> None:
+    schema = create_schema("test", {"o": dy.Object()})
+    assert schema.is_valid(
+        pl.DataFrame({"o": [CustomObject(a=1, b="foo"), CustomObject(a=2, b="bar")]})
+    )
+
+
+@pytest.mark.parametrize(
+    ("column", "dtype", "is_valid"),
+    [
+        (
+            dy.Object(),
+            pl.Object(),
+            True,
+        ),
+        (
+            dy.Object(),
+            object(),
+            False,
+        ),
+    ],
+)
+def test_validate_dtype(column: Column, dtype: pl.DataType, is_valid: bool) -> None:
+    assert column.validate_dtype(dtype) == is_valid
+
+
+def test_pyarrow_dtype_raises() -> None:
+    column = dy.Object()
+    with pytest.raises(
+        NotImplementedError, match="PyArrow column cannot have 'Object' type."
+    ):
+        column.pyarrow_dtype
+
+
+def test_sampling_raises() -> None:
+    column = dy.Object()
+    with pytest.raises(
+        NotImplementedError,
+        match="Random data sampling not implemented for 'Object' type.",
+    ):
+        column.sample(generator=Generator(), n=10)

{dataframely-1.1.0 → dataframely-1.2.1}/tests/columns/test_sql_schema.py

@@ -145,3 +145,11 @@ def test_raise_for_struct_column(dialect: sa.Dialect) -> None:
         NotImplementedError, match="SQL column cannot have 'Struct' type."
     ):
         dy.Struct({"a": dy.String()}).sqlalchemy_dtype(dialect)
+
+
+@pytest.mark.parametrize("dialect", [MSDialect_pyodbc(), PGDialect_psycopg2()])
+def test_raise_for_object_column(dialect: sa.Dialect) -> None:
+    with pytest.raises(
+        NotImplementedError, match="SQL column cannot have 'Object' type."
+    ):
+        dy.Object().sqlalchemy_dtype(dialect)

{dataframely-1.1.0 → dataframely-1.2.1}/tests/schema/test_validate.py

@@ -116,3 +116,12 @@ def test_success_multi_row_strip_cast(
     )
     assert_frame_equal(actual, expected)
     assert MySchema.is_valid(df, cast=True)
+
+
+@pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame])
+def test_group_rule_on_nulls(df_type: type[pl.DataFrame] | type[pl.LazyFrame]) -> None:
+    # The schema is violated because we have multiple "b" values for the same "a" value
+    df = df_type({"a": [None, None], "b": [1, 2]})
+    with pytest.raises(RuleValidationError):
+        MyComplexSchema.validate(df, cast=True)
+    assert not MyComplexSchema.is_valid(df, cast=True)

{dataframely-1.1.0 → dataframely-1.2.1}/tests/test_typing.py

@@ -9,7 +9,7 @@
 import datetime
 import decimal
 import functools
-from typing import Any
+from typing import Any, NotRequired, TypedDict
 
 import polars as pl
 import pytest
@@ -61,7 +61,21 @@ class MySecondSchema(dy.Schema):
     b = dy.Integer()
 
 
-class MyCollection(dy.Collection):
+class SamplingTypeFirst(TypedDict):
+    a: NotRequired[int]
+
+
+class SamplingTypeSecond(TypedDict):
+    a: NotRequired[int]
+    b: NotRequired[int]
+
+
+class SamplingType(TypedDict):
+    first: NotRequired[SamplingTypeFirst]
+    second: NotRequired[SamplingTypeSecond]
+
+
+class MyCollection(dy.Collection[SamplingType]):
     first: dy.LazyFrame[MyFirstSchema]
     second: dy.LazyFrame[MySecondSchema]
 
@@ -73,6 +87,12 @@ def test_collection_filter_return_value() -> None:
     assert len(failure["third"]) == 0  # type: ignore[misc]
 
 
+def test_collection_concat() -> None:
+    c1 = MyCollection.create_empty()
+    c2 = MyCollection.create_empty()
+    dy.concat_collection_members([c1, c2])
+
+
 # ------------------------------------------------------------------------------------ #
 #                                       ITER ROWS                                      #
 # ------------------------------------------------------------------------------------ #