fix: Properly type-check non-existent Schema attributes (#224)

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: borchero <22455425+borchero@users.noreply.github.com>
Co-authored-by: Oliver Borchert <oliver.borchert@quantco.com>

Author: Copilot

dataframely/_base_schema.py

@@ -8,7 +8,7 @@
 from abc import ABCMeta
 from copy import copy
 from dataclasses import dataclass, field
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import polars as pl
 
@@ -164,12 +164,16 @@ def __new__(
 
         return cls
 
-    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
+    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:
@@ -199,9 +203,9 @@ def _get_metadata(source: dict[str, Any]) -> Metadata:
     def __repr__(cls) -> str:
         parts = [f'[Schema "{cls.__name__}"]']
         parts.append(textwrap.indent("Columns:", prefix=" " * 2))
-        for name, col in cls.columns().items():
+        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():
+        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))

dataframely/_pydantic.py

@@ -20,12 +20,12 @@
 def _dict_to_df(schema_type: type[BaseSchema], data: dict) -> pl.DataFrame:
     return pl.from_dict(
         data,
-        schema=schema_type.to_polars_schema(),
+        schema=schema_type.to_polars_schema(),  # type: ignore[attr-defined]
     )
 
 
 def _validate_df_schema(schema_type: type[_S], df: pl.DataFrame) -> DataFrame[_S]:
-    if not schema_type.is_valid(df):
+    if not schema_type.is_valid(df):  # type: ignore[attr-defined]
         raise ValueError("DataFrame violates schema")
     return df  # type: ignore
 

dataframely/filter_result.py

@@ -319,7 +319,7 @@ def _sink(
         backend.sink_failure_info(
             lf=self._lf,
             serialized_rules=json.dumps(self._rule_columns),
-            serialized_schema=self.schema.serialize(),
+            serialized_schema=self.schema.serialize(),  # type: ignore[attr-defined]
             **kwargs,
         )
 
@@ -333,7 +333,7 @@ def _write(
         backend.write_failure_info(
             df=self._df,
             serialized_rules=json.dumps(self._rule_columns),
-            serialized_schema=self.schema.serialize(),
+            serialized_schema=self.schema.serialize(),  # type: ignore[attr-defined]
             **kwargs,
         )
 

tests/test_typing.py

@@ -80,3 +80,16 @@ def test_collection_concat() -> None:
     c1 = MyCollection.create_empty()
     c2 = MyCollection.create_empty()
     dy.concat_collection_members([c1, c2])
+
+
+# ------------------------------------------------------------------------------------ #
+#                                   ATTRIBUTE ACCESS                                   #
+# ------------------------------------------------------------------------------------ #
+
+
+def test_non_existent_column_access() -> None:
+    Schema.non_existing_col  # type: ignore[attr-defined]
+
+
+def test_valid_column_access() -> None:
+    Schema.a  # Should pass type checking