dataframely 1.0.0__cp311-abi3-win_amd64.whl

dataframely/__init__.py

@@ -0,0 +1,92 @@
+# 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 ._base_collection import CollectionMember
+from ._filter import filter
+from ._rule import rule
+from ._typing import DataFrame, LazyFrame
+from .collection import Collection
+from .columns import (
+    Any,
+    Bool,
+    Column,
+    Date,
+    Datetime,
+    Decimal,
+    Duration,
+    Enum,
+    Float,
+    Float32,
+    Float64,
+    Int8,
+    Int16,
+    Int32,
+    Int64,
+    Integer,
+    List,
+    String,
+    Struct,
+    Time,
+    UInt8,
+    UInt16,
+    UInt32,
+    UInt64,
+)
+from .config import Config
+from .failure import FailureInfo
+from .functional import (
+    concat_collection_members,
+    filter_relationship_one_to_at_least_one,
+    filter_relationship_one_to_one,
+)
+from .schema import Schema
+
+__all__ = [
+    "random",
+    "filter",
+    "rule",
+    "DataFrame",
+    "LazyFrame",
+    "Collection",
+    "CollectionMember",
+    "Config",
+    "FailureInfo",
+    "concat_collection_members",
+    "filter_relationship_one_to_at_least_one",
+    "filter_relationship_one_to_one",
+    "Schema",
+    "Any",
+    "Bool",
+    "Column",
+    "Date",
+    "Datetime",
+    "Decimal",
+    "Duration",
+    "Time",
+    "Enum",
+    "Float",
+    "Float32",
+    "Float64",
+    "Int8",
+    "Int16",
+    "Int32",
+    "Int64",
+    "Integer",
+    "UInt8",
+    "UInt16",
+    "UInt32",
+    "UInt64",
+    "String",
+    "Struct",
+    "List",
+]

dataframely/_base_collection.py

@@ -0,0 +1,285 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+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
+
+import polars as pl
+
+from ._filter import Filter
+from ._typing import LazyFrame as TypedLazyFrame
+from .exc import AnnotationImplementationError, ImplementationError
+from .schema import Schema
+
+_MEMBER_ATTR = "__dataframely_members__"
+_FILTER_ATTR = "__dataframely_filters__"
+
+
+@dataclass(kw_only=True)
+class CollectionMember:
+    """An annotation class that configures different behavior for a collection member.
+
+    Members:
+        ignored_in_filters: Indicates that a member should be ignored in the
+            ``@dy.filter`` methods of a collection. This also affects the computation
+            of the shared primary key in the collection.
+
+    Example:
+        .. code:: python
+
+            class MyCollection(dy.Collection):
+                a: dy.LazyFrame[MySchema1]
+                b: dy.LazyFrame[MySchema2]
+
+                ignored_member: Annotated[
+                    dy.LazyFrame[MySchema3],
+                    dy.CollectionMember(ignored_in_filters=True)
+                ]
+
+                @dy.filter
+                def my_filter(self) -> pl.DataFrame:
+                    return self.a.join(self.b, on="shared_key")
+    """
+
+    #: Whether the member should be ignored in the filter method.
+    ignored_in_filters: bool = False
+
+
+# --------------------------------------- UTILS -------------------------------------- #
+
+
+def _common_primary_keys(columns: Iterable[type[Schema]]) -> set[str]:
+    return set.intersection(*[set(schema.primary_keys()) for schema in columns])
+
+
+# ------------------------------------------------------------------------------------ #
+#                                    COLLECTION META                                   #
+# ------------------------------------------------------------------------------------ #
+
+
+@dataclass
+class MemberInfo(CollectionMember):
+    """Information about a member of a collection."""
+
+    #: The schema of the member.
+    schema: type[Schema]
+    #: Whether the member is optional.
+    is_optional: bool
+
+
+@dataclass
+class Metadata:
+    """Utility class to gather members and filters associated with a collection."""
+
+    members: dict[str, MemberInfo] = field(default_factory=dict)
+    filters: dict[str, Filter] = field(default_factory=dict)
+
+    def update(self, other: Self):
+        self.members.update(other.members)
+        self.filters.update(other.filters)
+
+
+class CollectionMeta(ABCMeta):
+    def __new__(
+        mcs,  # noqa: N804
+        name: str,
+        bases: tuple[type[object], ...],
+        namespace: dict[str, Any],
+        *args: Any,
+        **kwargs: Any,
+    ):
+        result = Metadata()
+        for base in bases:
+            result.update(mcs._get_metadata_recursively(base))
+        result.update(mcs._get_metadata(namespace))
+        namespace[_MEMBER_ATTR] = result.members
+        namespace[_FILTER_ATTR] = result.filters
+
+        # We now have all necessary information about filters and members. We want to
+        # check some preconditions to not run into issues later...
+
+        non_ignored_member_schemas = [
+            m.schema for m in result.members.values() if not m.ignored_in_filters
+        ]
+
+        # 1) Check that there are overlapping primary keys that allow the application
+        # of filters.
+        if len(non_ignored_member_schemas) > 0 and len(result.filters) > 0:
+            if len(_common_primary_keys(non_ignored_member_schemas)) == 0:
+                raise ImplementationError(
+                    "Members of a collection must have an overlapping primary key "
+                    "but did not find any."
+                )
+
+        # 2) Check that filter names do not overlap with any column or rule names
+        if len(result.members) > 0:
+            taken = set.union(
+                *(
+                    set(member.schema.column_names())
+                    for member in result.members.values()
+                ),
+                *(
+                    set(member.schema._validation_rules())
+                    for member in result.members.values()
+                ),
+            )
+            intersection = taken & set(result.filters)
+            if len(intersection) > 0:
+                raise ImplementationError(
+                    "Filters defined on the collection must not be named the same as any "
+                    "column or rule in any of the member frames but found "
+                    f"{len(intersection)} such filters: {sorted(intersection)}."
+                )
+
+        return super().__new__(mcs, name, bases, namespace, *args, **kwargs)
+
+    @staticmethod
+    def _get_metadata_recursively(kls: type[object]) -> Metadata:
+        result = Metadata()
+        for base in kls.__bases__:
+            result.update(CollectionMeta._get_metadata_recursively(base))
+        result.update(CollectionMeta._get_metadata(kls.__dict__))  # type: ignore
+        return result
+
+    @staticmethod
+    def _get_metadata(source: dict[str, Any]) -> Metadata:
+        result = Metadata()
+
+        # 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,
+                    )
+                else:
+                    # Some other unknown annotation
+                    raise AnnotationImplementationError(attr, kls)
+
+        # Get all filters by traversing the source
+        for attr, value in {
+            k: v for k, v in source.items() if not k.startswith("__")
+        }.items():
+            if isinstance(value, Filter):
+                result.filters[attr] = value
+
+        return result
+
+
+class BaseCollection(metaclass=CollectionMeta):
+    """Internal utility abstraction to reference collections without introducing
+    cyclical dependencies."""
+
+    @classmethod
+    def members(cls) -> dict[str, MemberInfo]:
+        """Information about the members of the collection."""
+        return getattr(cls, _MEMBER_ATTR)
+
+    @classmethod
+    def member_schemas(cls) -> dict[str, type[Schema]]:
+        """The schemas of all members of the collection."""
+        return {name: member.schema for name, member in cls.members().items()}
+
+    @classmethod
+    def required_members(cls) -> set[str]:
+        """The names of all required members of the collection."""
+        return {
+            name for name, member in cls.members().items() if not member.is_optional
+        }
+
+    @classmethod
+    def optional_members(cls) -> set[str]:
+        """The names of all optional members of the collection."""
+        return {name for name, member in cls.members().items() if member.is_optional}
+
+    @classmethod
+    def ignored_members(cls) -> set[str]:
+        """The names of all members of the collection that are ignored in filters."""
+        return {
+            name for name, member in cls.members().items() if member.ignored_in_filters
+        }
+
+    @classmethod
+    def non_ignored_members(cls) -> set[str]:
+        """The names of all members of the collection that are not ignored in filters
+        (default)."""
+        return {
+            name
+            for name, member in cls.members().items()
+            if not member.ignored_in_filters
+        }
+
+    @classmethod
+    def common_primary_keys(cls) -> list[str]:
+        """The primary keys shared by non ignored members of the collection."""
+        return sorted(
+            _common_primary_keys(
+                [
+                    member.schema
+                    for member in cls.members().values()
+                    if not member.ignored_in_filters
+                ]
+            )
+        )
+
+    @classmethod
+    def _filters(cls) -> dict[str, Filter[Self]]:
+        return getattr(cls, _FILTER_ATTR)
+
+    def to_dict(self) -> dict[str, pl.LazyFrame]:
+        """Return a dictionary representation of this collection."""
+        return {
+            member: getattr(self, member)
+            for member in self.member_schemas()
+            if getattr(self, member) is not None
+        }

dataframely/_base_schema.py

@@ -0,0 +1,186 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+from abc import ABCMeta
+from copy import copy
+from dataclasses import dataclass, field
+from typing import Any, Self
+
+import polars as pl
+
+from ._rule import GroupRule, Rule, with_evaluation_rules
+from .columns import Column
+from .exc import ImplementationError, RuleImplementationError
+
+_COLUMN_ATTR = "__dataframely_columns__"
+_RULE_ATTR = "__dataframely_rules__"
+
+# --------------------------------------- UTILS -------------------------------------- #
+
+
+def _build_rules(
+    custom: dict[str, Rule], columns: dict[str, Column]
+) -> 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_keys = _primary_keys(columns)
+    if len(primary_keys) > 0:
+        rules["primary_key"] = Rule(~pl.struct(primary_keys).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)
+
+    return rules
+
+
+def _primary_keys(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, Rule] = field(default_factory=dict)
+
+    def update(self, other: Self):
+        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,
+    ):
+        result = Metadata()
+        for base in bases:
+            result.update(mcs._get_metadata_recursively(base))
+        result.update(mcs._get_metadata(namespace))
+        namespace[_COLUMN_ATTR] = result.columns
+        namespace[_RULE_ATTR] = result.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(result.rules, result.columns).keys()) | set(
+            f"{col}|dtype" for col in result.columns
+        )
+        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 name, rule in result.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 '{name}' has been implemented "
+                        f"incorrectly. It references {len(missing_columns)} columns "
+                        f"which are not in the schema: {missing_list}."
+                    )
+
+        # 3) Assuming that non-custom rules are implemented correctly, we check that all
+        # custom rules are _also_ implemented correctly by evaluating rules on an
+        # empty data frame and checking for the evaluated dtypes.
+        if len(result.rules) > 0:
+            lf_empty = pl.LazyFrame(
+                schema={name: col.dtype for name, col in result.columns.items()}
+            )
+            # NOTE: For some reason, `polars` does not yield correct dtypes when calling
+            #  `collect_schema()`
+            schema = with_evaluation_rules(lf_empty, result.rules).collect().schema
+            for name, rule in result.rules.items():
+                dtype = schema[name]
+                if not isinstance(dtype, pl.Boolean):
+                    raise RuleImplementationError(
+                        name, dtype, isinstance(rule, GroupRule)
+                    )
+
+        return super().__new__(mcs, name, bases, namespace, *args, **kwargs)
+
+    @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):
+                if not value.alias:
+                    value.alias = attr
+                result.columns[value.alias] = value
+            if isinstance(value, Rule):
+                # 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
+
+
+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."""
+        return getattr(cls, _COLUMN_ATTR)
+
+    @classmethod
+    def primary_keys(cls) -> list[str]:
+        """The primary key columns in this schema (possibly empty)."""
+        return _primary_keys(cls.columns())
+
+    @classmethod
+    def _validation_rules(cls) -> dict[str, Rule]:
+        return _build_rules(cls._schema_validation_rules(), cls.columns())
+
+    @classmethod
+    def _schema_validation_rules(cls) -> dict[str, Rule]:
+        return getattr(cls, _RULE_ATTR)

dataframely/_compat.py

@@ -0,0 +1,39 @@
+# 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):
+        self.module = module
+
+    def __getattr__(self, name: str) -> Any:
+        raise ValueError(f"Module '{self.module}' is not installed.")
+
+
+# ------------------------------------ SQLALCHEMY ------------------------------------ #
+
+try:
+    import sqlalchemy as sa
+    import sqlalchemy.dialects.mssql as sa_mssql
+    from sqlalchemy.sql.type_api import TypeEngine as sa_TypeEngine
+except ImportError:  # pragma: no cover
+    sa = _DummyModule("sqlalchemy")  # type: ignore
+    sa_mssql = _DummyModule("sqlalchemy")  # type: ignore
+
+    class sa_TypeEngine:  # type: ignore # noqa: N801
+        pass
+
+
+# -------------------------------------- PYARROW ------------------------------------- #
+
+try:
+    import pyarrow as pa
+except ImportError:  # pragma: no cover
+    pa = _DummyModule("pyarrow")
+
+# ------------------------------------------------------------------------------------ #
+
+__all__ = ["sa", "sa_mssql", "sa_TypeEngine", "pa"]

dataframely/_extre.pyi

@@ -0,0 +1,54 @@
+from typing import Literal, overload
+
+def 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 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 sample(
+    regex: str,
+    n: Literal[None] = None,
+    max_repetitions: int = 16,
+    seed: int | None = None,
+) -> str: ...

dataframely/_filter.py

@@ -0,0 +1,42 @@
+# 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]):
+        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.
+    """
+
+    def decorator(validation_fn: Callable[[C], pl.LazyFrame]) -> Filter[C]:
+        return Filter(logic=validation_fn)
+
+    return decorator