feat: Allow for multiple column level checks with custom naming (#19)

Co-authored-by: Oliver Borchert <[email protected]>

Author: heshamdar

dataframely/columns/_base.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections import Counter
 from collections.abc import Callable
 from typing import Any
 
@@ -31,7 +32,12 @@ def __init__(
         *,
         nullable: bool | None = None,
         primary_key: bool = False,
-        check: Callable[[pl.Expr], pl.Expr] | None = None,
+        check: (
+            Callable[[pl.Expr], pl.Expr]
+            | list[Callable[[pl.Expr], pl.Expr]]
+            | dict[str, Callable[[pl.Expr], pl.Expr]]
+            | None
+        ) = None,
         alias: str | None = None,
         metadata: dict[str, Any] | None = None,
     ):
@@ -43,8 +49,17 @@ def __init__(
                 is not specified.
             primary_key: Whether this column is part of the primary key of the schema.
                 If ``True``, ``nullable`` is automatically set to ``False``.
-            check: A custom check to run for this column. Must return a non-aggregated
-                boolean expression.
+            check: A custom rule or multiple rules to run for this column. This can be:
+                - A single callable that returns a non-aggregated boolean expression.
+                The name of the rule is derived from the callable name, or defaults to
+                "check" for lambdas.
+                - A list of callables, where each callable returns a non-aggregated
+                boolean expression. The name of the rule is derived from the callable
+                name, or defaults to "check" for lambdas. Where multiple rules result
+                in the same name, the suffix __i is appended to the name.
+                - A dictionary mapping rule names to callables, where each callable
+                returns a non-aggregated boolean expression.
+                All rule names provided here are given the prefix "check_".
             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
@@ -104,10 +119,59 @@ def validation_rules(self, expr: pl.Expr) -> dict[str, pl.Expr]:
         result = {}
         if not self.nullable:
             result["nullability"] = expr.is_not_null()
+
         if self.check is not None:
-            result["check"] = self.check(expr)
+            if isinstance(self.check, dict):
+                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]
+                )
+                # Get unique names for rules from callables
+                rule_names = self._derive_check_rule_names(list_of_rules)
+                for rule_name, rule_callable in zip(rule_names, list_of_rules):
+                    result[rule_name] = rule_callable(expr)
+
         return result
 
+    def _derive_check_rule_names(
+        self, rules: list[Callable[[pl.Expr], pl.Expr]]
+    ) -> list[str]:
+        """Generate unique names for rule callables.
+
+        For callables with the same name, appends a suffix __i where i is the index
+        of occurrence (starting from 0), but only if there are duplicates.
+
+        Args:
+            rules: List of rule callables.
+
+        Returns:
+            List of unique names corresponding to the rule callables.
+        """
+        base_names = [
+            f"check__{rule.__name__}" if rule.__name__ != "<lambda>" else "check"
+            for rule in rules
+        ]
+
+        # Count occurrences using Counter
+        name_counts = Counter(base_names)
+
+        # Append suffixes to names that are duplicated
+        final_names = []
+        duplicate_counter: dict[str, int] = {
+            name: 0 for name in name_counts if name_counts[name] > 1
+        }
+        for name in base_names:
+            if name_counts[name] > 1:
+                postfix = duplicate_counter[name]
+                final_names.append(f"{name}__{postfix}")
+                duplicate_counter[name] += 1
+            else:
+                final_names.append(name)
+
+        return final_names
+
     # -------------------------------------- SQL ------------------------------------- #
 
     def sqlalchemy_column(self, name: str, dialect: sa.Dialect) -> sa.Column:

dataframely/columns/any.py

@@ -25,14 +25,28 @@ class Any(Column):
     def __init__(
         self,
         *,
-        check: Callable[[pl.Expr], pl.Expr] | None = None,
+        check: (
+            Callable[[pl.Expr], pl.Expr]
+            | list[Callable[[pl.Expr], pl.Expr]]
+            | dict[str, Callable[[pl.Expr], pl.Expr]]
+            | None
+        ) = None,
         alias: str | None = None,
         metadata: dict[str, Any] | None = None,
     ):
         """
         Args:
-            check: A custom check to run for this column. Must return a non-aggregated
-                boolean expression.
+            check: A custom rule or multiple rules to run for this column. This can be:
+                - A single callable that returns a non-aggregated boolean expression.
+                The name of the rule is derived from the callable name, or defaults to
+                "check" for lambdas.
+                - A list of callables, where each callable returns a non-aggregated
+                boolean expression. The name of the rule is derived from the callable
+                name, or defaults to "check" for lambdas. Where multiple rules result
+                in the same name, the suffix __i is appended to the name.
+                - A dictionary mapping rule names to callables, where each callable
+                returns a non-aggregated boolean expression.
+                All rule names provided here are given the prefix "check_".
             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

dataframely/columns/array.py

@@ -28,7 +28,12 @@ def __init__(
         # polars doesn't yet support grouping by arrays,
         # see https://github.com/pola-rs/polars/issues/22574
         primary_key: Literal[False] = False,
-        check: Callable[[pl.Expr], pl.Expr] | None = None,
+        check: (
+            Callable[[pl.Expr], pl.Expr]
+            | list[Callable[[pl.Expr], pl.Expr]]
+            | dict[str, Callable[[pl.Expr], pl.Expr]]
+            | None
+        ) = None,
         alias: str | None = None,
         metadata: dict[str, Any] | None = None,
     ):
@@ -39,8 +44,17 @@ def __init__(
             nullable: Whether this column may contain null values.
             primary_key: Whether this column is part of the primary key of the schema.
                 Not yet supported for the Array type.
-            check: A custom check to run for this column. Must return a non-aggregated
-                boolean expression.
+            check: A custom rule or multiple rules to run for this column. This can be:
+                - A single callable that returns a non-aggregated boolean expression.
+                The name of the rule is derived from the callable name, or defaults to
+                "check" for lambdas.
+                - A list of callables, where each callable returns a non-aggregated
+                boolean expression. The name of the rule is derived from the callable
+                name, or defaults to "check" for lambdas. Where multiple rules result
+                in the same name, the suffix __i is appended to the name.
+                - A dictionary mapping rule names to callables, where each callable
+                returns a non-aggregated boolean expression.
+                All rule names provided here are given the prefix "check_".
             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

dataframely/columns/datetime.py

@@ -39,7 +39,12 @@ def __init__(
         max: dt.date | None = None,
         max_exclusive: dt.date | None = None,
         resolution: str | None = None,
-        check: Callable[[pl.Expr], pl.Expr] | None = None,
+        check: (
+            Callable[[pl.Expr], pl.Expr]
+            | list[Callable[[pl.Expr], pl.Expr]]
+            | dict[str, Callable[[pl.Expr], pl.Expr]]
+            | None
+        ) = None,
         alias: str | None = None,
         metadata: dict[str, Any] | None = None,
     ):
@@ -61,8 +66,17 @@ def __init__(
                 formatting language used by :mod:`polars` datetime ``round`` method.
                 For example, a value ``1mo`` expects all dates to be on the first of the
                 month. Note that this setting does *not* affect the storage resolution.
-            check: A custom check to run for this column. Must return a non-aggregated
-                boolean expression.
+            check: A custom rule or multiple rules to run for this column. This can be:
+                - A single callable that returns a non-aggregated boolean expression.
+                The name of the rule is derived from the callable name, or defaults to
+                "check" for lambdas.
+                - A list of callables, where each callable returns a non-aggregated
+                boolean expression. The name of the rule is derived from the callable
+                name, or defaults to "check" for lambdas. Where multiple rules result
+                in the same name, the suffix __i is appended to the name.
+                - A dictionary mapping rule names to callables, where each callable
+                returns a non-aggregated boolean expression.
+                All rule names provided here are given the prefix "check_".
             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
@@ -152,7 +166,12 @@ def __init__(
         max: dt.time | None = None,
         max_exclusive: dt.time | None = None,
         resolution: str | None = None,
-        check: Callable[[pl.Expr], pl.Expr] | None = None,
+        check: (
+            Callable[[pl.Expr], pl.Expr]
+            | list[Callable[[pl.Expr], pl.Expr]]
+            | dict[str, Callable[[pl.Expr], pl.Expr]]
+            | None
+        ) = None,
         alias: str | None = None,
         metadata: dict[str, Any] | None = None,
     ):
@@ -174,8 +193,17 @@ def __init__(
                 formatting language used by :mod:`polars` datetime ``round`` method.
                 For example, a value ``1h`` expects all times to be full hours. Note
                 that this setting does *not* affect the storage resolution.
-            check: A custom check to run for this column. Must return a non-aggregated
-                boolean expression.
+            check: A custom rule or multiple rules to run for this column. This can be:
+                - A single callable that returns a non-aggregated boolean expression.
+                The name of the rule is derived from the callable name, or defaults to
+                "check" for lambdas.
+                - A list of callables, where each callable returns a non-aggregated
+                boolean expression. The name of the rule is derived from the callable
+                name, or defaults to "check" for lambdas. Where multiple rules result
+                in the same name, the suffix __i is appended to the name.
+                - A dictionary mapping rule names to callables, where each callable
+                returns a non-aggregated boolean expression.
+                All rule names provided here are given the prefix "check_".
             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
@@ -271,7 +299,12 @@ def __init__(
         max: dt.datetime | None = None,
         max_exclusive: dt.datetime | None = None,
         resolution: str | None = None,
-        check: Callable[[pl.Expr], pl.Expr] | None = None,
+        check: (
+            Callable[[pl.Expr], pl.Expr]
+            | list[Callable[[pl.Expr], pl.Expr]]
+            | dict[str, Callable[[pl.Expr], pl.Expr]]
+            | None
+        ) = None,
         alias: str | None = None,
         metadata: dict[str, Any] | None = None,
     ):
@@ -293,8 +326,17 @@ def __init__(
                 the formatting language used by :mod:`polars` datetime ``round`` method.
                 For example, a value ``1h`` expects all datetimes to be full hours. Note
                 that this setting does *not* affect the storage resolution.
-            check: A custom check to run for this column. Must return a non-aggregated
-                boolean expression.
+            check: A custom rule or multiple rules to run for this column. This can be:
+                - A single callable that returns a non-aggregated boolean expression.
+                The name of the rule is derived from the callable name, or defaults to
+                "check" for lambdas.
+                - A list of callables, where each callable returns a non-aggregated
+                boolean expression. The name of the rule is derived from the callable
+                name, or defaults to "check" for lambdas. Where multiple rules result
+                in the same name, the suffix __i is appended to the name.
+                - A dictionary mapping rule names to callables, where each callable
+                returns a non-aggregated boolean expression.
+                All rule names provided here are given the prefix "check_".
             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
@@ -380,7 +422,12 @@ def __init__(
         max: dt.timedelta | None = None,
         max_exclusive: dt.timedelta | None = None,
         resolution: str | None = None,
-        check: Callable[[pl.Expr], pl.Expr] | None = None,
+        check: (
+            Callable[[pl.Expr], pl.Expr]
+            | list[Callable[[pl.Expr], pl.Expr]]
+            | dict[str, Callable[[pl.Expr], pl.Expr]]
+            | None
+        ) = None,
         alias: str | None = None,
         metadata: dict[str, Any] | None = None,
     ):
@@ -402,8 +449,17 @@ def __init__(
                 the formatting language used by :mod:`polars` datetime ``round`` method.
                 For example, a value ``1h`` expects all durations to be full hours. Note
                 that this setting does *not* affect the storage resolution.
-            check: A custom check to run for this column. Must return a non-aggregated
-                boolean expression.
+            check: A custom rule or multiple rules to run for this column. This can be:
+                - A single callable that returns a non-aggregated boolean expression.
+                The name of the rule is derived from the callable name, or defaults to
+                "check" for lambdas.
+                - A list of callables, where each callable returns a non-aggregated
+                boolean expression. The name of the rule is derived from the callable
+                name, or defaults to "check" for lambdas. Where multiple rules result
+                in the same name, the suffix __i is appended to the name.
+                - A dictionary mapping rule names to callables, where each callable
+                returns a non-aggregated boolean expression.
+                All rule names provided here are given the prefix "check_".
             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

dataframely/columns/decimal.py

@@ -33,7 +33,12 @@ def __init__(
         min_exclusive: decimal.Decimal | None = None,
         max: decimal.Decimal | None = None,
         max_exclusive: decimal.Decimal | None = None,
-        check: Callable[[pl.Expr], pl.Expr] | None = None,
+        check: (
+            Callable[[pl.Expr], pl.Expr]
+            | list[Callable[[pl.Expr], pl.Expr]]
+            | dict[str, Callable[[pl.Expr], pl.Expr]]
+            | None
+        ) = None,
         alias: str | None = None,
         metadata: dict[str, Any] | None = None,
     ):
@@ -53,8 +58,17 @@ def __init__(
             max: The maximum value for decimals in this column (inclusive).
             max_exclusive: Like ``max`` but exclusive. May not be specified if ``max``
                 is specified and vice versa.
-            check: A custom check to run for this column. Must return a non-aggregated
-                boolean expression.
+            check: A custom rule or multiple rules to run for this column. This can be:
+                - A single callable that returns a non-aggregated boolean expression.
+                The name of the rule is derived from the callable name, or defaults to
+                "check" for lambdas.
+                - A list of callables, where each callable returns a non-aggregated
+                boolean expression. The name of the rule is derived from the callable
+                name, or defaults to "check" for lambdas. Where multiple rules result
+                in the same name, the suffix __i is appended to the name.
+                - A dictionary mapping rule names to callables, where each callable
+                returns a non-aggregated boolean expression.
+                All rule names provided here are given the prefix "check_".
             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