accpet only None in Union

Author: gabriel

dataframely/collection/_base.py

@@ -244,15 +244,17 @@ def _derive_member_info(
                 attr, annotation_args[0], annotation_args[1]
             )
         elif origin == typing.Union:
-            # Happy path: optional member
+            # Happy path: optional member (e.g. dy.LazyFrame[Schema] | None)
             union_args = get_args(type_annotation)
             if len(union_args) != 2:
                 raise AnnotationImplementationError(attr, type_annotation)
-            if not any(get_origin(arg) is None for arg in union_args):
+            # Check that exactly one arg is None (type(None) is NoneType)
+            if not any(arg is type(None) for arg in union_args):
                 raise AnnotationImplementationError(attr, type_annotation)
 
-            not_none_args = [arg for arg in union_args if get_origin(arg) is not None]
-            if len(not_none_args) == 0:
+            # Get the non-None type
+            not_none_args = [arg for arg in union_args if arg is not type(None)]
+            if len(not_none_args) != 1:
                 raise AnnotationImplementationError(attr, type_annotation)
 
             frame_origin = get_origin(not_none_args[0])

tests/collection/test_implementation.py

@@ -113,7 +113,7 @@ def test_annotation_only_none_failure() -> None:
 
 
 def test_annotation_invalid_type_failure() -> None:
-    """First argument of union must be a LazyFrame."""
+    """First argument of union must be a LazyFrame or DataFrame."""
     with pytest.raises(AnnotationImplementationError):
         create_collection_raw(
             "test",
@@ -123,6 +123,39 @@ def test_annotation_invalid_type_failure() -> None:
         )
 
 
+def test_annotation_invalid_generic_type_in_union() -> None:
+    """Union with generic type that's not LazyFrame/DataFrame should fail."""
+    with pytest.raises(AnnotationImplementationError):
+        create_collection_raw(
+            "test",
+            {
+                "first": list[int] | None,
+            },
+        )
+
+
+def test_annotation_union_frame_with_non_none_type() -> None:
+    """Union of DataFrame with non-None type should fail."""
+    with pytest.raises(AnnotationImplementationError):
+        create_collection_raw(
+            "test",
+            {
+                "first": dy.DataFrame[MyTestSchema] | int,
+            },
+        )
+
+
+def test_annotation_invalid_standalone_generic() -> None:
+    """Standalone generic type that's not LazyFrame/DataFrame should fail."""
+    with pytest.raises(AnnotationImplementationError):
+        create_collection_raw(
+            "test",
+            {
+                "first": list[int],
+            },
+        )
+
+
 def test_explicit_annotation_type_failure_no_frame_type() -> None:
     """First argument of the annotated union must be a LazyFrame."""
     with pytest.raises(AnnotationImplementationError):