dataframely 2.2.0__cp310-abi3-win_amd64.whl

dataframely/__init__.py

@@ -0,0 +1,111 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+import importlib.metadata
+import warnings
+
+try:
+    __version__ = importlib.metadata.version(__name__)
+except importlib.metadata.PackageNotFoundError as e:  # pragma: no cover
+    warnings.warn(f"Could not determine version of {__name__}\n{e!s}", stacklevel=2)
+    __version__ = "unknown"
+
+from . import random
+from ._filter import filter
+from ._rule import rule
+from ._typing import DataFrame, LazyFrame, Validation
+from .collection import (
+    Collection,
+    CollectionMember,
+    deserialize_collection,
+    read_parquet_metadata_collection,
+)
+from .columns import (
+    Any,
+    Array,
+    Binary,
+    Bool,
+    Categorical,
+    Column,
+    Date,
+    Datetime,
+    Decimal,
+    Duration,
+    Enum,
+    Float,
+    Float32,
+    Float64,
+    Int8,
+    Int16,
+    Int32,
+    Int64,
+    Integer,
+    List,
+    Object,
+    String,
+    Struct,
+    Time,
+    UInt8,
+    UInt16,
+    UInt32,
+    UInt64,
+)
+from .config import Config
+from .exc import DeserializationError
+from .filter_result import FailureInfo
+from .functional import (
+    concat_collection_members,
+    require_relationship_one_to_at_least_one,
+    require_relationship_one_to_one,
+)
+from .schema import Schema, deserialize_schema, read_parquet_metadata_schema
+
+__all__ = [
+    "random",
+    "filter",
+    "rule",
+    "DataFrame",
+    "LazyFrame",
+    "Collection",
+    "CollectionMember",
+    "deserialize_collection",
+    "Config",
+    "FailureInfo",
+    "concat_collection_members",
+    "require_relationship_one_to_at_least_one",
+    "require_relationship_one_to_one",
+    "Schema",
+    "deserialize_schema",
+    "read_parquet_metadata_schema",
+    "read_parquet_metadata_collection",
+    "Any",
+    "Binary",
+    "Bool",
+    "Categorical",
+    "Column",
+    "Date",
+    "Datetime",
+    "Decimal",
+    "Duration",
+    "Time",
+    "Enum",
+    "Float",
+    "Float32",
+    "Float64",
+    "Int8",
+    "Int16",
+    "Int32",
+    "Int64",
+    "Integer",
+    "UInt8",
+    "UInt16",
+    "UInt32",
+    "UInt64",
+    "String",
+    "Struct",
+    "List",
+    "Array",
+    "Object",
+    "Validation",
+    "DeserializationError",
+]

dataframely/_base_schema.py

@@ -0,0 +1,247 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import sys
+import textwrap
+from abc import ABCMeta
+from copy import copy
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+import polars as pl
+
+from ._rule import DtypeCastRule, GroupRule, Rule, RuleFactory
+from .columns import Column
+from .exc import ImplementationError
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+
+_COLUMN_ATTR = "__dataframely_columns__"
+_RULE_ATTR = "__dataframely_rules__"
+
+ORIGINAL_COLUMN_PREFIX = "__DATAFRAMELY_ORIGINAL__"
+
+# --------------------------------------- UTILS -------------------------------------- #
+
+
+def _build_rules(
+    custom: dict[str, Rule], columns: dict[str, Column], *, with_cast: bool
+) -> dict[str, Rule]:
+    # NOTE: Copy here to prevent in-place modification of the custom rules
+    rules: dict[str, Rule] = copy(custom)
+
+    # Add primary key validation to the list of rules if applicable
+    primary_key = _primary_key(columns)
+    if len(primary_key) > 0:
+        rules["primary_key"] = Rule(~pl.struct(primary_key).is_duplicated())
+
+    # Add column-specific rules
+    column_rules = {
+        f"{col_name}|{rule_name}": Rule(expr)
+        for col_name, column in columns.items()
+        for rule_name, expr in column.validation_rules(pl.col(col_name)).items()
+    }
+    rules.update(column_rules)
+
+    # Add casting rules if requested. Here, we can simply check whether the nullability
+    # property of a column changes due to lenient dtype casting. Whenever casting fails,
+    # the value is set to `null`, mismatching the previous nullability.
+    # NOTE: This check assumes that both the original and cast column are present in the
+    #  data frame.
+    if with_cast:
+        casting_rules = {
+            f"{col_name}|dtype": DtypeCastRule(
+                pl.col(col_name).is_null()
+                == pl.col(f"{ORIGINAL_COLUMN_PREFIX}{col_name}").is_null()
+            )
+            for col_name in columns
+        }
+        rules.update(casting_rules)
+
+    return rules
+
+
+def _primary_key(columns: dict[str, Column]) -> list[str]:
+    return list(k for k, col in columns.items() if col.primary_key)
+
+
+# ------------------------------------------------------------------------------------ #
+#                                      SCHEMA META                                     #
+# ------------------------------------------------------------------------------------ #
+
+
+@dataclass
+class Metadata:
+    """Utility class to gather columns and rules associated with a schema."""
+
+    columns: dict[str, Column] = field(default_factory=dict)
+    rules: dict[str, RuleFactory] = field(default_factory=dict)
+
+    def update(self, other: Self) -> None:
+        self.columns.update(other.columns)
+        self.rules.update(other.rules)
+
+
+class SchemaMeta(ABCMeta):
+    def __new__(
+        mcs,  # noqa: N804
+        name: str,
+        bases: tuple[type[object], ...],
+        namespace: dict[str, Any],
+        *args: Any,
+        **kwargs: Any,
+    ) -> SchemaMeta:
+        result = Metadata()
+        for base in bases:
+            result.update(mcs._get_metadata_recursively(base))
+        result.update(mcs._get_metadata(namespace))
+        namespace[_COLUMN_ATTR] = result.columns
+        cls = super().__new__(mcs, name, bases, namespace, *args, **kwargs)
+
+        # Assign rules retroactively as we only encounter rule factories in the result
+        rules = {name: factory.make(cls) for name, factory in result.rules.items()}
+        setattr(cls, _RULE_ATTR, rules)
+
+        # At this point, we already know all columns and custom rules. We want to run
+        # some checks...
+
+        # 1) Check that the column names clash with none of the rule names. To this end,
+        # we assume that users cast dtypes, i.e. additional rules for dtype casting
+        # are also checked.
+        all_column_names = set(result.columns)
+        all_rule_names = set(_build_rules(rules, result.columns, with_cast=True))
+        common_names = all_column_names & all_rule_names
+        if len(common_names) > 0:
+            common_list = ", ".join(sorted(f"'{col}'" for col in common_names))
+            raise ImplementationError(
+                "Rules and columns must not be named equally but found "
+                f"{len(common_names)} overlaps: {common_list}."
+            )
+
+        # 2) Check that the columns referenced in the group rules exist.
+        for rule_name, rule in rules.items():
+            if isinstance(rule, GroupRule):
+                missing_columns = set(rule.group_columns) - set(result.columns)
+                if len(missing_columns) > 0:
+                    missing_list = ", ".join(
+                        sorted(f"'{col}'" for col in missing_columns)
+                    )
+                    raise ImplementationError(
+                        f"Group validation rule '{rule_name}' has been implemented "
+                        f"incorrectly. It references {len(missing_columns)} columns "
+                        f"which are not in the schema: {missing_list}."
+                    )
+
+        # 3) Check that all members are non-pathological (i.e., user errors).
+        for attr, value in namespace.items():
+            if attr.startswith("__"):
+                continue
+
+            # Check for tuple of column (commonly caused by trailing comma)
+            if (
+                isinstance(value, tuple)
+                and len(value) > 0
+                and isinstance(value[0], Column)
+            ):
+                raise TypeError(
+                    f"Column '{attr}' is defined as a tuple of dy.Column. "
+                    f"Did you accidentally add a trailing comma?"
+                )
+
+            # Check for column type instead of instance (e.g., dy.Float64 instead of dy.Float64())
+            if isinstance(value, type) and issubclass(value, Column):
+                raise TypeError(
+                    f"Column '{attr}' is a type, not an instance. "
+                    f"Schema members must be of type Column not type[Column]. "
+                    f"Did you forget to add parentheses?"
+                )
+
+        return cls
+
+    if not TYPE_CHECKING:
+        # Only define __getattribute__ at runtime to allow type checkers to properly
+        # validate attribute access. When TYPE_CHECKING is True, type checkers will use
+        # the default metaclass behavior which correctly identifies non-existent attributes.
+        def __getattribute__(cls, name: str) -> Any:
+            val = super().__getattribute__(name)
+            # Dynamically set the name of the column if it is a `Column` instance.
+            if isinstance(val, Column):
+                val._name = val.alias or name
+            return val
+
+    @staticmethod
+    def _get_metadata_recursively(kls: type[object]) -> Metadata:
+        result = Metadata()
+        for base in kls.__bases__:
+            result.update(SchemaMeta._get_metadata_recursively(base))
+        result.update(SchemaMeta._get_metadata(kls.__dict__))  # type: ignore
+        return result
+
+    @staticmethod
+    def _get_metadata(source: dict[str, Any]) -> Metadata:
+        result = Metadata()
+        for attr, value in {
+            k: v for k, v in source.items() if not k.startswith("__")
+        }.items():
+            if isinstance(value, Column):
+                result.columns[value.alias or attr] = value
+            if isinstance(value, RuleFactory):
+                # We must ensure that custom rules do not clash with internal rules.
+                if attr == "primary_key":
+                    raise ImplementationError(
+                        "Custom validation rule must not be named `primary_key`."
+                    )
+                result.rules[attr] = value
+        return result
+
+    def __repr__(cls) -> str:
+        parts = [f'[Schema "{cls.__name__}"]']
+        parts.append(textwrap.indent("Columns:", prefix=" " * 2))
+        for name, col in cls.columns().items():  # type: ignore[attr-defined]
+            parts.append(textwrap.indent(f'- "{name}": {col!r}', prefix=" " * 4))
+        if validation_rules := cls._schema_validation_rules():  # type: ignore[attr-defined]
+            parts.append(textwrap.indent("Rules:", prefix=" " * 2))
+            for name, rule in validation_rules.items():
+                parts.append(textwrap.indent(f'- "{name}": {rule!r}', prefix=" " * 4))
+        parts.append("")  # Add line break at the end
+        return "\n".join(parts)
+
+
+class BaseSchema(metaclass=SchemaMeta):
+    """Internal utility abstraction to reference schemas without introducing cyclical
+    dependencies."""
+
+    @classmethod
+    def column_names(cls) -> list[str]:
+        """The column names of this schema."""
+        return list(getattr(cls, _COLUMN_ATTR).keys())
+
+    @classmethod
+    def columns(cls) -> dict[str, Column]:
+        """The column definitions of this schema."""
+        columns: dict[str, Column] = getattr(cls, _COLUMN_ATTR)
+        for name in columns.keys():
+            # Dynamically set the name of the columns.
+            columns[name]._name = name
+        return columns
+
+    @classmethod
+    def primary_key(cls) -> list[str]:
+        """The primary key columns in this schema (possibly empty)."""
+        return _primary_key(cls.columns())
+
+    @classmethod
+    def _validation_rules(cls, *, with_cast: bool) -> dict[str, Rule]:
+        return _build_rules(
+            cls._schema_validation_rules(), cls.columns(), with_cast=with_cast
+        )
+
+    @classmethod
+    def _schema_validation_rules(cls) -> dict[str, Rule]:
+        return getattr(cls, _RULE_ATTR)

dataframely/_compat.py

@@ -0,0 +1,86 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+
+from typing import Any
+
+
+class _DummyModule:  # pragma: no cover
+    def __init__(self, module: str) -> None:
+        self.module = module
+
+    def __getattr__(self, name: str) -> Any:
+        raise ValueError(f"Module '{self.module}' is not installed.")
+
+
+# ------------------------------------ DELTALAKE ------------------------------------- #
+
+try:
+    import deltalake
+    from deltalake import DeltaTable
+except ImportError:
+    deltalake = _DummyModule("deltalake")  # type: ignore
+
+    class DeltaTable:  # type: ignore # noqa: N801
+        pass
+
+# ------------------------------------ SQLALCHEMY ------------------------------------ #
+
+try:
+    import sqlalchemy as sa
+    import sqlalchemy.dialects.mssql as sa_mssql
+    from sqlalchemy import Dialect
+    from sqlalchemy.dialects.mssql.pyodbc import MSDialect_pyodbc
+    from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2
+    from sqlalchemy.sql.type_api import TypeEngine as sa_TypeEngine
+except ImportError:
+    sa = _DummyModule("sqlalchemy")  # type: ignore
+    sa_mssql = _DummyModule("sqlalchemy")  # type: ignore
+
+    class sa_TypeEngine:  # type: ignore # noqa: N801
+        pass
+
+    class MSDialect_pyodbc:  # type: ignore # noqa: N801
+        pass
+
+    class PGDialect_psycopg2:  # type: ignore # noqa: N801
+        pass
+
+    class Dialect:  # type: ignore # noqa: N801
+        pass
+
+# -------------------------------------- PYARROW ------------------------------------- #
+
+try:
+    import pyarrow as pa
+except ImportError:
+    pa = _DummyModule("pyarrow")
+
+
+# -------------------------------------- PYDANTIC ------------------------------------ #
+
+try:
+    import pydantic
+except ImportError:
+    pydantic = _DummyModule("pydantic")  # type: ignore
+
+try:
+    from pydantic_core import core_schema as pydantic_core_schema
+except ImportError:
+    pydantic_core_schema = _DummyModule("pydantic_core_schema")  # type: ignore
+
+# ------------------------------------------------------------------------------------ #
+
+__all__ = [
+    "deltalake",
+    "DeltaTable",
+    "Dialect",
+    "MSDialect_pyodbc",
+    "pa",
+    "PGDialect_psycopg2",
+    "pydantic_core_schema",
+    "pydantic",
+    "sa_mssql",
+    "sa_TypeEngine",
+    "sa",
+]

dataframely/_deprecation.py

@@ -0,0 +1,27 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+import os
+from collections.abc import Callable
+from functools import wraps
+
+TRUTHY_VALUES = ["1", "true"]
+
+
+def skip_if(env: str) -> Callable:
+    """Decorator to skip warnings based on environment variable.
+
+    If the environment variable is equivalent to any of TRUTHY_VALUES, the wrapped
+    function is skipped.
+    """
+
+    def decorator(fun: Callable) -> Callable:
+        @wraps(fun)
+        def wrapper() -> None:
+            if os.getenv(env, "").lower() in TRUTHY_VALUES:
+                return
+            fun()
+
+        return wrapper
+
+    return decorator

dataframely/_filter.py

@@ -0,0 +1,48 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+from collections.abc import Callable
+from typing import Generic, TypeVar
+
+import polars as pl
+
+C = TypeVar("C")
+
+
+class Filter(Generic[C]):
+    """Internal class representing logic for filtering members of a collection."""
+
+    def __init__(self, logic: Callable[[C], pl.LazyFrame]) -> None:
+        self.logic = logic
+
+
+def filter() -> Callable[[Callable[[C], pl.LazyFrame]], Filter[C]]:
+    """Mark a function as filters for rows in the members of a collection.
+
+    The name of the function will be used as the name of the filter. The name must not
+    clash with the name of any column in the member schemas or rules defined on the
+    member schemas.
+
+    A filter receives a collection as input and must return a data frame like the
+    following:
+
+    - The columns must be a superset of the common primary keys across all members.
+    - The rows must provide the primary keys which ought to be *kept* across the
+      members. The filter results in the removal of rows which are lost as the result
+      of inner-joining members onto the return value of this function.
+
+    Attention:
+        Make sure to provide unique combinations of the primary keys or the filters
+        might introduce duplicate rows.
+
+    Attention:
+        The filter logic should return a lazy frame with a static computational graph.
+        Other implementations using arbitrary python logic works for filtering and
+        validation, but may lead to wrong results in Collection comparisons
+        and (de-)serialization.
+    """
+
+    def decorator(validation_fn: Callable[[C], pl.LazyFrame]) -> Filter[C]:
+        return Filter(logic=validation_fn)
+
+    return decorator

dataframely/_match_to_schema.py

@@ -0,0 +1,89 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+from typing import Literal
+
+import polars as pl
+
+from ._base_schema import ORIGINAL_COLUMN_PREFIX, BaseSchema
+from .columns import Column
+from .exc import SchemaError
+
+
+def match_to_schema(
+    lf: pl.LazyFrame,
+    target: type[BaseSchema],
+    *,
+    casting: Literal["none", "lenient", "strict"],
+) -> pl.LazyFrame:
+    """Ensure that a lazy frame contains the columns of the schema with the dtypes
+    specified by the schema."""
+
+    def cast_none(lf: pl.LazyFrame, schema: pl.Schema) -> pl.LazyFrame:
+        _validate_columns_exist(schema, target)
+        _validate_dtypes(schema, target)
+        return lf.select(target.column_names())
+
+    def cast_lenient(lf: pl.LazyFrame, schema: pl.Schema) -> pl.LazyFrame:
+        _validate_columns_exist(schema, target)
+        # NOTE: We keep around the original columns for failure objects and
+        #  to evaluate whether casting is successful.
+        return lf.select(
+            pl.col(target.column_names()).name.prefix(ORIGINAL_COLUMN_PREFIX),
+            *[
+                pl.col(name).pipe(_cast_if_required, schema[name], column, strict=False)
+                for name, column in target.columns().items()
+            ],
+        )
+
+    def cast_strict(lf: pl.LazyFrame, schema: pl.Schema) -> pl.LazyFrame:
+        _validate_columns_exist(schema, target)
+        return lf.select(
+            pl.col(name).pipe(_cast_if_required, schema[name], column)
+            for name, column in target.columns().items()
+        )
+
+    match casting:
+        case "none":
+            return lf.pipe_with_schema(cast_none)
+        case "lenient":
+            return lf.pipe_with_schema(cast_lenient)
+        case "strict":
+            return lf.pipe_with_schema(cast_strict)
+
+
+# ------------------------------------------------------------------------------------ #
+
+
+def _validate_columns_exist(actual: pl.Schema, target: type[BaseSchema]) -> None:
+    actual_columns = set(actual.keys())
+    target_columns = set(target.column_names())
+    if missing := target_columns - actual_columns:
+        raise SchemaError(
+            f"{len(missing)} missing columns for schema '{target.__name__}': "
+            + ", ".join(f"'{c}'" for c in sorted(missing))
+        )
+
+
+def _validate_dtypes(actual: pl.Schema, target: type[BaseSchema]) -> None:
+    failures = {
+        name: column
+        for name, column in target.columns().items()
+        if not column.validate_dtype(actual[name])
+    }
+    if failures:
+        raise SchemaError(
+            f"{len(failures)} columns with invalid dtype for schema '{target.__name__}': "
+            + "\n".join(
+                f" - '{name}', got: {actual[name]}, expected: {column}"
+                for name, column in failures.items()
+            )
+        )
+
+
+def _cast_if_required(
+    expr: pl.Expr, current_dtype: pl.DataType, column: Column, *, strict: bool = True
+) -> pl.Expr:
+    if column.validate_dtype(current_dtype):
+        return expr
+    return expr.cast(column.dtype, strict=strict)

dataframely/_native.pyi

@@ -0,0 +1,67 @@
+from typing import overload
+
+def format_rule_failures(failures: list[tuple[str, int]]) -> str:
+    """
+    Format rule failures with the same logic that produces validation errors from the
+    polars plugin.
+
+    Args:
+        failures: The name of the failures and their counts. This should only include
+            failures with a count of at least 1.
+
+    Returns:
+        The formatted rule failures.
+    """
+
+def regex_matching_string_length(regex: str) -> tuple[int, int | None]:
+    """
+    Compute the minimum and maximum length (if available) of strings matching a regular expression.
+
+    Args:
+        regex: The regular expression to analyze. The regular expression must not
+            contain any lookaround operators.
+
+    Returns:
+        A tuple of the minimum of maximum length of the matching strings. While the minimum
+        length is guaranteed to be available, the maximum length may be `None` if `regex`
+        matches strings of potentially infinite length (e.g. due to the use of `+` or `*`).
+
+    Raises:
+        ValueError: If the regex cannot be parsed or analyzed.
+    """
+
+@overload
+def regex_sample(
+    regex: str, n: int, max_repetitions: int = 16, seed: int | None = None
+) -> list[str]:
+    """
+    Sample a random (set of) string(s) matching the provided regular expression.
+
+    Args:
+        regex: The regular expression generated strings must match. The regular
+            expression must not contain any lookaround operators.
+        n: The number of random strings to generate or `None` if a single one should
+            be generated.
+        max_repetitions: The maximum number of repetitions for `+` and `*`
+            quantifiers.
+        seed: The seed to use for the random sampling procedure.
+
+    Returns:
+        A single randomly generated string if `n is None` or a list of randomly
+        generated strings if `n` is an integer.
+
+    Raises:
+        ValueError: If the regex cannot be parsed.
+
+    Attention:
+        Using wildcards (i.e. `.`) really means _any_ valid Unicode character.
+        Consider using more precise regular expressions if this is undesired.
+    """
+
+@overload
+def regex_sample(
+    regex: str,
+    n: None = None,
+    max_repetitions: int = 16,
+    seed: int | None = None,
+) -> str: ...