Allow null values during dataframe validation for optional enum fields (

#44)

src/patito/validators.py

@@ -120,6 +120,8 @@ def _find_errors(  # noqa: C901
         # Test for when only specific values are accepted
         if "enum" in column_properties:
             permissible_values = set(column_properties["enum"])
+            if column_name in schema.nullable_columns:
+                permissible_values.add(None)
             actual_values = set(dataframe[column_name].unique())
             impermissible_values = actual_values - permissible_values
             if impermissible_values:

tests/test_validators.py

@@ -1,5 +1,6 @@
 """Tests for the patito.validators module."""
 import enum
+import sys
 from datetime import date, datetime
 from typing import Optional
 
@@ -478,3 +479,34 @@ class Pair(pt.Model):
                 ]
             )
         )
+
+
+def test_optional_enum():
+    """It should handle optional enums correctly."""
+
+    class OptionalEnumModel(pt.Model):
+        # Old type annotation syntax
+        optional_enum: Optional[Literal["A", "B"]]
+
+    df = pl.DataFrame({"optional_enum": ["A", "B", None]})
+    OptionalEnumModel.validate(df)
+
+
+@pytest.mark.skipif(
+    sys.version_info <= (3, 10),
+    reason="Using | as a type union operator is only supported from python 3.10.",
+)
+def test_optional_pipe_operator():
+    class OptionalEnumModel(pt.Model):
+        # Old type annotation syntax
+        optional_enum_1: Optional[Literal["A", "B"]]
+        # New type annotation syntax
+        optional_enum_2: Literal["A", "B"] | None  # pyright: ignore
+
+    df = pl.DataFrame(
+        {
+            "optional_enum_1": ["A", "B", None],
+            "optional_enum_2": ["A", "B", None],
+        }
+    )
+    OptionalEnumModel.validate(df)