dataframely 1.10.0__tar.gz → 1.12.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataframely
-Version: 1.10.0
+Version: 1.12.0
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11

{dataframely-1.10.0 → dataframely-1.12.0}/dataframely/__init__.py

@@ -23,6 +23,7 @@ from .collection import (
 from .columns import (
     Any,
     Array,
+    Binary,
     Bool,
     Categorical,
     Column,
@@ -77,6 +78,7 @@ __all__ = [
     "read_parquet_metadata_schema",
     "read_parquet_metadata_collection",
     "Any",
+    "Binary",
     "Bool",
     "Categorical",
     "Column",

{dataframely-1.10.0 → dataframely-1.12.0}/dataframely/columns/__init__.py

@@ -5,6 +5,7 @@ from ._base import Column
 from ._registry import column_from_dict
 from .any import Any
 from .array import Array
+from .binary import Binary
 from .bool import Bool
 from .categorical import Categorical
 from .datetime import Date, Datetime, Duration, Time
@@ -22,6 +23,7 @@ __all__ = [
     "column_from_dict",
     "Any",
     "Array",
+    "Binary",
     "Bool",
     "Categorical",
     "Date",

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

@@ -7,7 +7,7 @@ import inspect
 import sys
 from abc import ABC, abstractmethod
 from collections import Counter
-from collections.abc import Callable
+from collections.abc import Callable, Mapping, Sequence
 from typing import Any, TypeAlias, cast
 
 import polars as pl
@@ -27,8 +27,8 @@ else:
 
 Check: TypeAlias = (
     Callable[[pl.Expr], pl.Expr]
-    | list[Callable[[pl.Expr], pl.Expr]]
-    | dict[str, Callable[[pl.Expr], pl.Expr]]
+    | Sequence[Callable[[pl.Expr], pl.Expr]]
+    | Mapping[str, Callable[[pl.Expr], pl.Expr]]
 )
 
 # ------------------------------------------------------------------------------------ #
@@ -141,12 +141,14 @@ class Column(ABC):
             result["nullability"] = expr.is_not_null()
 
         if self.check is not None:
-            if isinstance(self.check, dict):
+            if isinstance(self.check, Mapping):
                 for rule_name, rule_callable in self.check.items():
                     result[f"check__{rule_name}"] = rule_callable(expr)
             else:
                 list_of_rules = (
-                    self.check if isinstance(self.check, list) else [self.check]
+                    list(self.check)
+                    if isinstance(self.check, Sequence)
+                    else [self.check]
                 )
                 # Get unique names for rules from callables
                 rule_names = self._derive_check_rule_names(list_of_rules)
@@ -378,6 +380,16 @@ class Column(ABC):
     ) -> bool:
         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 -------------------------------- #
@@ -401,6 +413,10 @@ 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):
@@ -423,9 +439,9 @@ def _check_to_expr(check: Check | None, expr: pl.Expr) -> Any | None:
     match check:
         case None:
             return None
-        case list():
+        case Sequence():
             return [c(expr) for c in check]
-        case dict():
+        case Mapping():
             return {key: c(expr) for key, c in check.items()}
         case _ if callable(check):
             return check(expr)

dataframely-1.12.0/dataframely/columns/binary.py

@@ -0,0 +1,38 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import polars as pl
+
+from dataframely._compat import pa, sa, sa_TypeEngine
+from dataframely.random import Generator
+
+from ._base import Column
+from ._registry import register
+
+
+@register
+class Binary(Column):
+    """A column of binary values."""
+
+    @property
+    def dtype(self) -> pl.DataType:
+        return pl.Binary()
+
+    def sqlalchemy_dtype(self, dialect: sa.Dialect) -> sa_TypeEngine:
+        if dialect.name == "mssql":
+            return sa.VARBINARY()
+        return sa.LargeBinary()
+
+    @property
+    def pyarrow_dtype(self) -> pa.DataType:
+        return pa.large_binary()
+
+    def _sample_unchecked(self, generator: Generator, n: int) -> pl.Series:
+        return generator.sample_binary(
+            n,
+            min_bytes=0,
+            max_bytes=32,
+            null_probability=self._null_probability,
+        )

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

@@ -3,7 +3,9 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+import enum
+from collections.abc import Iterable
+from inspect import isclass
 from typing import Any
 
 import polars as pl
@@ -22,7 +24,7 @@ class Enum(Column):
 
     def __init__(
         self,
-        categories: Sequence[str],
+        categories: pl.Series | Iterable[str] | type[enum.Enum],
         *,
         nullable: bool | None = None,
         primary_key: bool = False,
@@ -32,7 +34,8 @@ class Enum(Column):
     ):
         """
         Args:
-            categories: The list of valid categories for the enum.
+            categories: The set of valid categories for the enum, or an existing Python
+                string-valued enum.
             nullable: Whether this column may contain null values.
                 Explicitly set `nullable=True` if you want your column to be nullable.
                 In a future release, `nullable=False` will be the default if `nullable`
@@ -63,7 +66,13 @@ class Enum(Column):
             alias=alias,
             metadata=metadata,
         )
-        self.categories = list(categories)
+        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
 
     @property
     def dtype(self) -> pl.DataType:
@@ -72,7 +81,7 @@ class Enum(Column):
     def validate_dtype(self, dtype: PolarsDataType) -> bool:
         if not isinstance(dtype, pl.Enum):
             return False
-        return self.categories == dtype.categories.to_list()
+        return self.categories.equals(dtype.categories)
 
     def sqlalchemy_dtype(self, dialect: sa.Dialect) -> sa_TypeEngine:
         category_lengths = [len(c) for c in self.categories]
@@ -92,5 +101,7 @@ class Enum(Column):
 
     def _sample_unchecked(self, generator: Generator, n: int) -> pl.Series:
         return generator.sample_choice(
-            n, choices=self.categories, null_probability=self._null_probability
+            n,
+            choices=self.categories.to_list(),
+            null_probability=self._null_probability,
         ).cast(self.dtype)

{dataframely-1.10.0 → dataframely-1.12.0}/dataframely/random.py

@@ -148,6 +148,31 @@ class Generator:
             pl.Series(samples, dtype=pl.String), null_probability
         )
 
+    def sample_binary(
+        self,
+        n: int = 1,
+        *,
+        min_bytes: int,
+        max_bytes: int,
+        null_probability: float = 0.0,
+    ) -> pl.Series:
+        """Sample a list of binary values in the specified length range.
+
+        Args:
+            n: The number of binary values to sample.
+            min_bytes: The minimum number of bytes for each value.
+            max_bytes: The maximum number of bytes for each value.
+            null_probability: The probability of an element being ``null``.
+
+        Returns:
+            A series with ``n`` elements of dtype ``Binary``.
+        """
+        lengths = self.numpy_generator.integers(min_bytes, max_bytes + 1, size=n)
+        samples = [self.numpy_generator.bytes(length) for length in lengths]
+        return self._apply_null_mask(
+            pl.Series(samples, dtype=pl.Binary), null_probability
+        )
+
     def sample_choice(
         self,
         n: int = 1,

{dataframely-1.10.0 → dataframely-1.12.0}/dataframely/testing/const.py

@@ -22,6 +22,7 @@ COLUMN_TYPES: list[type[dc.Column]] = [
     dc.UInt64,
     dc.String,
     dc.Categorical,
+    dc.Binary,
 ]
 INTEGER_COLUMN_TYPES: list[type[dc.Column]] = [
     dc.Integer,

{dataframely-1.10.0 → dataframely-1.12.0}/pyproject.toml

@@ -25,7 +25,7 @@ description = "A declarative, polars-native data frame validation library"
 name = "dataframely"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.10.0"
+version = "1.12.0"
 
 [project.optional-dependencies]
 deltalake = ["deltalake"]

dataframely-1.12.0/tests/column_types/test_binary.py

@@ -0,0 +1,24 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+import polars as pl
+import pytest
+
+import dataframely as dy
+from dataframely.columns import Column
+
+
+class BinarySchema(dy.Schema):
+    a = dy.Binary()
+
+
+@pytest.mark.parametrize(
+    ("column", "dtype", "is_valid"),
+    [
+        (dy.Binary(), pl.Binary(), True),
+        (dy.Binary(), pl.String(), False),
+        (dy.Binary(), pl.Null(), False),
+    ],
+)
+def test_validate_dtype(column: Column, dtype: pl.DataType, is_valid: bool) -> None:
+    assert column.validate_dtype(dtype) == is_valid

{dataframely-1.10.0 → dataframely-1.12.0}/tests/column_types/test_enum.py

@@ -1,6 +1,8 @@
 # Copyright (c) QuantCo 2025-2025
 # SPDX-License-Identifier: BSD-3-Clause
-
+import enum
+from collections.abc import Iterable
+from enum import Enum
 from typing import Any
 
 import polars as pl
@@ -61,3 +63,48 @@ def test_different_sequences(type1: type, type2: type) -> None:
     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)
+
+
+def test_enum_of_enum_136() -> None:
+    class Categories(str, Enum):
+        a = "a"
+        b = "b"
+
+    assert pl.Enum(Categories) == dy.Enum(Categories).dtype
+
+
+def test_enum_of_series() -> None:
+    categories = pl.Series(["a", "b"])
+    assert pl.Enum(categories) == dy.Enum(categories).dtype
+
+
+def test_enum_of_iterable() -> None:
+    categories = (x for x in ["a", "b"])
+    assert pl.Enum(["a", "b"]) == dy.Enum(categories).dtype
+
+
+@pytest.mark.parametrize(
+    "categories1",
+    [
+        ["a", "b"],
+        ("a", "b"),
+        pl.Series(["a", "b"]),
+        Enum("Categories", {"a": "a", "b": "b"}),
+    ],
+)
+@pytest.mark.parametrize(
+    "categories2",
+    [
+        ["a", "b"],
+        ("a", "b"),
+        pl.Series(["a", "b"]),
+        Enum("Categories", {"a": "a", "b": "b"}),
+    ],
+)
+def test_sequences_and_enums(
+    categories1: pl.Series | Iterable[str] | type[enum.Enum],
+    categories2: pl.Series | Iterable[str] | type[enum.Enum],
+) -> None:
+    S = create_schema("test", {"x": dy.Enum(categories1)})
+    df = pl.DataFrame({"x": pl.Series(["a", "b"], dtype=pl.Enum(categories2))})
+    S.validate(df)

{dataframely-1.10.0 → dataframely-1.12.0}/tests/columns/test_check.py

@@ -12,6 +12,22 @@ class CheckSchema(dy.Schema):
     b = dy.String(min_length=3, check=lambda col: col.str.contains("x"))
 
 
+def test_check_argument_covariant() -> None:
+    # The interesting part of this test is mypy accepting it statically,
+    # not the runtime behavior.
+    def check_all_or_none(expr: pl.Expr) -> pl.Expr:
+        return (expr == "all").all() | (expr != "all").all()
+
+    check_dict = {"all_or_none": check_all_or_none}
+
+    class CheckDictSchema(dy.Schema):
+        column = dy.String(check=check_dict)
+
+    df = pl.DataFrame({"column": ["all", "all", "all"]})
+    _, failures = CheckDictSchema.filter(df)
+    assert failures.counts() == {}
+
+
 def test_check() -> None:
     df = pl.DataFrame({"a": [7, 3, 15], "b": ["abc", "xyz", "x"]})
     _, failures = CheckSchema.filter(df)

{dataframely-1.10.0 → dataframely-1.12.0}/tests/columns/test_matches.py

@@ -58,7 +58,19 @@ import dataframely as dy
             dy.Datetime(time_zone=dt.timezone(dt.timedelta(hours=0))),
             True,
         ),
+        (dy.Enum(["a", "b"]), dy.Enum(["a", "b"]), True),
+        (dy.Enum(["a", "b"]), dy.Enum(["a", "b", "c"]), False),
     ],
 )
 def test_matches(lhs: dy.Column, rhs: dy.Column, expected: bool) -> None:
     assert lhs.matches(rhs, expr=pl.element()) == expected
+
+
+def test_matches_enum_attribute_type_mismatch() -> None:
+    # Comparison should fail if the `other` column has
+    # a `category` member, but its dtype is not `pl.Series`
+    col1 = dy.Enum(["a", "b"])
+    col2 = dy.Enum(["a", "b"])
+    col2.categories = "this_is_not_a_series"  # type: ignore
+
+    assert not col1.matches(col2, pl.element())

{dataframely-1.10.0 → dataframely-1.12.0}/tests/columns/test_sql_schema.py

@@ -15,6 +15,7 @@ pytestmark = pytest.mark.with_optionals
     ("column", "datatype"),
     [
         (dy.Any(), "SQL_VARIANT"),
+        (dy.Binary(), "VARBINARY(max)"),
         (dy.Bool(), "BIT"),
         (dy.Date(), "DATE"),
         (dy.Datetime(), "DATETIME2(6)"),
@@ -61,6 +62,7 @@ def test_mssql_datatype(column: Column, datatype: str) -> None:
 @pytest.mark.parametrize(
     ("column", "datatype"),
     [
+        (dy.Binary(), "BYTEA"),
         (dy.Bool(), "BOOLEAN"),
         (dy.Date(), "DATE"),
         (dy.Datetime(), "TIMESTAMP WITHOUT TIME ZONE"),

{dataframely-1.10.0 → dataframely-1.12.0}/tests/test_random.py

@@ -40,6 +40,7 @@ def test_seeding_nonconstant() -> None:
         lambda generator, n: generator.sample_float(n, min=0, max=5),
         lambda generator, n: generator.sample_string(n, regex="[abc]"),
         lambda generator, n: generator.sample_choice(n, choices=[1, 2, 3]),
+        lambda generator, n: generator.sample_binary(n, min_bytes=1, max_bytes=10),
         lambda generator, n: generator.sample_time(n, min=dt.time(0, 0), max=None),
         lambda generator, n: generator.sample_date(
             n, min=dt.date(1970, 1, 1), max=None
@@ -75,6 +76,9 @@ def test_sample_correct_n(
         lambda generator, n, prob: generator.sample_choice(
             n, choices=[1, 2, 3], null_probability=prob
         ),
+        lambda generator, n, prob: generator.sample_binary(
+            n, min_bytes=1, max_bytes=10, null_probability=prob
+        ),
         lambda generator, n, prob: generator.sample_time(
             n, min=dt.time(0, 0), max=None, null_probability=prob
         ),
@@ -131,6 +135,16 @@ def test_sample_string(generator: Generator) -> None:
     assert (samples.str.len_bytes() == 2).all()
 
 
+def test_sample_binary(generator: Generator) -> None:
+    samples = generator.sample_binary(100, min_bytes=1, max_bytes=10)
+    assert (
+        samples.to_frame("s").select(pl.col("s").bin.size("b") >= 1).to_series().all()
+    )
+    assert (
+        samples.to_frame("s").select(pl.col("s").bin.size("b") <= 10).to_series().all()
+    )
+
+
 def test_sample_choice(generator: Generator) -> None:
     samples = generator.sample_choice(100_000, choices=[1, 2, 3])
     assert np.allclose(