feat: add way to configure attributes like nullable, cast, more conveniently using selectors

pyproject.toml

@@ -14,6 +14,8 @@ dependencies = ["narwhals>=1.0.0"]
 [project.optional-dependencies]
 pyperclip = ["pyperclip>=1.0.0"]
 
+[dependency-groups]
+typing = ["typing_extensions"]
 
 [build-system]
 requires = ["setuptools>=42", "wheel"]
@@ -27,3 +29,4 @@ profile = "black"
 
 [tool.mypy]
 plugins = ["checkedframe.mypy"]
+allow_redefinition = true

src/checkedframe/__init__.py

@@ -2,6 +2,7 @@
 from checkedframe import exceptions, selectors
 
 from ._checks import Check, col, lit
+from ._config import Config, apply_configs
 from ._core import Schema
 from ._dtypes import (
     Array,

src/checkedframe/_config.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, TypedDict, Union
+
+from .selectors import Selector, by_name
+
+if TYPE_CHECKING:
+    from typing_extensions import Unpack
+
+SelectorLike = Union[str, Iterable[str], Selector]
+
+
+class _PossibleConfigs(TypedDict):
+    nullable: bool
+    required: bool
+    cast: bool
+    allow_nan: bool
+    allow_inf: bool
+
+
+class Config:
+    def __init__(
+        self,
+        selector: SelectorLike,
+        **kwargs: Unpack[_PossibleConfigs],
+    ):
+        if not isinstance(selector, Selector):
+            actual_selector = by_name(selector)
+        else:
+            actual_selector = selector
+
+        self.selector = actual_selector
+        self.dct = kwargs
+
+
+# This just makes it easier to do isinstance checks
+class ConfigList:
+    def __init__(self, *args: Config):
+        self.args = args
+
+
+def apply_configs(*args: Config):
+    def decorator(cls):
+        cls.__private_checkedframe_config = ConfigList(*args)
+
+        return cls
+
+    return decorator

src/checkedframe/_core.py

@@ -8,7 +8,8 @@
 import narwhals.stable.v1 as nw
 import narwhals.stable.v1.typing as nwt
 
-from ._checks import Check, CheckInputType
+from ._checks import Check
+from ._config import ConfigList
 from ._dtypes import CastError, _Column, _nw_type_to_cf_type, _TypedColumn
 from ._utils import get_class_members
 from .exceptions import SchemaError
@@ -153,7 +154,7 @@ class InterrogationResult:
 
 
 def _private_interrogate(
-    schema: Schema, df: nwt.IntoDataFrameT, cast: bool
+    schema: Schema, df: nwt.IntoDataFrameT
 ) -> _PrivateInterrogationResult:
     nw_df = nw.from_native(df, eager_only=True)
     df_schema = nw_df.collect_schema()  # type: ignore[attribute]
@@ -201,7 +202,7 @@ def _private_interrogate(
         if actual_dtype == expected_col.to_narwhals():
             pass
         else:
-            if expected_col.cast or cast:
+            if expected_col.cast:
                 try:
                     nw_df = nw_df.with_columns(
                         _nw_type_to_cf_type(actual_dtype)._safe_cast(
@@ -351,9 +352,10 @@ def _private_interrogate(
 
 
 def _interrogate(
-    schema: Schema, df: nwt.IntoDataFrameT, cast: bool
+    schema: Schema,
+    df: nwt.IntoDataFrameT,
 ) -> InterrogationResult:
-    res = _private_interrogate(schema=schema, df=df, cast=cast)
+    res = _private_interrogate(schema=schema, df=df)
 
     return InterrogationResult(
         df=res.df.to_native(),
@@ -365,8 +367,8 @@ def _interrogate(
     )
 
 
-def _filter(schema: Schema, df: nwt.IntoDataFrameT, cast: bool) -> nwt.IntoDataFrameT:
-    res = _private_interrogate(schema=schema, df=df, cast=cast)
+def _filter(schema: Schema, df: nwt.IntoDataFrameT) -> nwt.IntoDataFrameT:
+    res = _private_interrogate(schema=schema, df=df)
 
     return res.df.filter(res.is_good).to_native()
 
@@ -430,8 +432,8 @@ def _wrap_err(e: str) -> str:
     return "\n".join(error_summary + output)
 
 
-def _validate(schema: Schema, df: nwt.IntoDataFrameT, cast: bool) -> nwt.IntoDataFrameT:
-    res = _private_interrogate(schema, df, cast)
+def _validate(schema: Schema, df: nwt.IntoDataFrameT) -> nwt.IntoDataFrameT:
+    res = _private_interrogate(schema, df)
 
     if not res.is_good.all():
         raise SchemaError(
@@ -569,33 +571,33 @@ def _parse_into_schema(cls) -> Schema:
                 else:
                     checks.append(val)
 
+            if isinstance(val, ConfigList):
+                for config in val.args:
+                    for c in config.selector(schema_dict):
+                        for k, v in config.dct.items():
+                            setattr(schema_dict[c], k, v)
+
         res = Schema(expected_schema=schema_dict, checks=checks)
         cls._schema = res
 
         return res
 
     @classmethod
-    def interrogate(
-        cls, df: nwt.IntoDataFrameT, cast: bool = False
-    ) -> InterrogationResult:
-        return _interrogate(cls._parse_into_schema(), df, cast)
+    def interrogate(cls, df: nwt.IntoDataFrameT) -> InterrogationResult:
+        return _interrogate(cls._parse_into_schema(), df)
 
-    def __interrogate(
-        self, df: nwt.IntoDataFrameT, cast: bool = False
-    ) -> InterrogationResult:
-        return _interrogate(self, df, cast)
+    def __interrogate(self, df: nwt.IntoDataFrameT) -> InterrogationResult:
+        return _interrogate(self, df)
 
     @classmethod
-    def validate(cls, df: nwt.IntoDataFrameT, cast: bool = False) -> nwt.IntoDataFrameT:
+    def validate(cls, df: nwt.IntoDataFrameT) -> nwt.IntoDataFrameT:
         """Validate the given DataFrame
 
         Parameters
         ----------
         df : nwt.IntoDataFrameT
             Any Narwhals-compatible DataFrame, see https://narwhals-dev.github.io/narwhals/
             for more information
-        cast : bool, optional
-            Whether to cast columns, by default False
 
         Returns
         -------
@@ -607,18 +609,14 @@ def validate(cls, df: nwt.IntoDataFrameT, cast: bool = False) -> nwt.IntoDataFra
         SchemaError
             If validation fails
         """
-        return _validate(cls._parse_into_schema(), df, cast)
+        return _validate(cls._parse_into_schema(), df)
 
-    def __validate(
-        self, df: nwt.IntoDataFrameT, cast: bool = False
-    ) -> nwt.IntoDataFrameT:
-        return _validate(self, df, cast)
+    def __validate(self, df: nwt.IntoDataFrameT) -> nwt.IntoDataFrameT:
+        return _validate(self, df)
 
     @classmethod
-    def filter(cls, df: nwt.IntoDataFrameT, cast: bool = False) -> nwt.IntoDataFrameT:
-        return _filter(cls._parse_into_schema(), df, cast)
+    def filter(cls, df: nwt.IntoDataFrameT) -> nwt.IntoDataFrameT:
+        return _filter(cls._parse_into_schema(), df)
 
-    def __filter(
-        self, df: nwt.IntoDataFrameT, cast: bool = False
-    ) -> nwt.IntoDataFrameT:
-        return _filter(self, df, cast)
+    def __filter(self, df: nwt.IntoDataFrameT) -> nwt.IntoDataFrameT:
+        return _filter(self, df)

src/checkedframe/_dtypes.py

@@ -34,14 +34,6 @@ class _BoundedDType(_DType):
     _max: int | float
 
 
-class _ColumnKwargs(TypedDict):
-    name: Optional[str]
-    nullable: bool
-    required: bool
-    cast: bool
-    checks: Optional[list[Check]]
-
-
 class _Column:
     """Represents a column in a DataFrame.
 

src/checkedframe/selectors.py

@@ -41,6 +41,14 @@ def __xor__(self, other: Selector) -> Selector:
             lambda col, dtype: self.condition(col, dtype) != other.condition(col, dtype)
         )
 
+    def exclude(self, other: str | Iterable[str] | Selector) -> Selector:
+        if not isinstance(other, Selector):
+            selector_other = by_name(other)
+        else:
+            selector_other = other
+
+        return self.__sub__(selector_other)
+
 
 def _flatten_str_iterable(lst: Iterable[str | Iterable[str]]) -> list[str]:
     res = []

tests/test_config.py

@@ -0,0 +1,24 @@
+import checkedframe as cf
+
+
+def test_config():
+    @cf.apply_configs(cf.Config(cf.selectors.float(), nullable=True))
+    class S(cf.Schema):
+        float1 = cf.Float64()
+        float2 = cf.Float64()
+        int1 = cf.Int64()
+
+    s = S._parse_into_schema().expected_schema
+
+    assert s["float1"].nullable
+    assert s["float2"].nullable
+    assert not s["int1"].nullable
+
+    class S2(S):
+        float3 = cf.Float64()
+
+    s2 = S2._parse_into_schema().expected_schema
+
+    assert s["float1"].nullable
+    assert s["float2"].nullable
+    assert s2["float3"].nullable

tests/test_selectors.py

@@ -72,3 +72,9 @@ def test_boolean():
 
 def test_numeric():
     assert set(cfs.numeric()(SCHEMA)) == set(["int8", "float32"])
+
+
+def test_exclude():
+    assert cfs.numeric().exclude("int8")(SCHEMA) == ["float32"]
+    assert cfs.numeric().exclude(cfs.by_name("int8"))(SCHEMA) == ["float32"]
+    assert cfs.numeric().exclude(["int8"])(SCHEMA) == ["float32"]