refine errors a bit

Author: johnnygreco

src/data_designer/config/utils/errors.py

@@ -10,4 +10,10 @@ class UserJinjaTemplateSyntaxError(DataDesignerError): ...
 class InvalidEnumValueError(DataDesignerError): ...
 
 
+class InvalidTypeUnionError(DataDesignerError): ...
+
+
+class InvalidDiscriminatorFieldError(DataDesignerError): ...
+
+
 class DatasetSampleDisplayError(DataDesignerError): ...

src/data_designer/config/utils/type_helpers.py

@@ -3,12 +3,12 @@
 
 from enum import Enum
 import inspect
-from typing import Any, Type, Union, get_args
+from typing import Any, Literal, Type, Union, get_args, get_origin
 
 from pydantic import BaseModel
 
 from .. import sampler_params
-from .errors import InvalidEnumValueError
+from .errors import InvalidDiscriminatorFieldError, InvalidEnumValueError, InvalidTypeUnionError
 
 
 class StrEnum(str, Enum):
@@ -43,9 +43,11 @@ def create_str_enum_from_discriminated_type_union(
     discriminator_field_values = []
     for model in type_union.__args__:
         if not issubclass(model, BaseModel):
-            raise ValueError(f"🛑 {model} is not a Pydantic model.")
+            raise InvalidTypeUnionError(f"🛑 {model} must be a subclass of pydantic.BaseModel.")
         if discriminator_field_name not in model.model_fields:
-            raise ValueError(f"🛑 {discriminator_field_name} is not a field of {model}.")
+            raise InvalidDiscriminatorFieldError(f"🛑 '{discriminator_field_name}' is not a field of {model}.")
+        if get_origin(model.model_fields[discriminator_field_name].annotation) is not Literal:
+            raise InvalidDiscriminatorFieldError(f"🛑 '{discriminator_field_name}' must be a Literal type.")
         discriminator_field_values.extend(get_args(model.model_fields[discriminator_field_name].annotation))
     return StrEnum(enum_name, {v.replace("-", "_").upper(): v for v in set(discriminator_field_values)})
 

tests/config/utils/test_type_helpers.py

@@ -7,7 +7,11 @@
 from pydantic import BaseModel
 import pytest
 
-from data_designer.config.utils.errors import InvalidEnumValueError
+from data_designer.config.utils.errors import (
+    InvalidDiscriminatorFieldError,
+    InvalidEnumValueError,
+    InvalidTypeUnionError,
+)
 from data_designer.config.utils.type_helpers import (
     SAMPLER_PARAMS,
     create_str_enum_from_discriminated_type_union,
@@ -94,16 +98,19 @@ class StubModelD(BaseModel):
 def test_create_str_enum_from_type_union_not_pydantic_model() -> None:
     type_union = Union[StubModelA, NotAModel]
 
-    with pytest.raises(ValueError, match="is not a Pydantic model"):
+    with pytest.raises(InvalidTypeUnionError, match="must be a subclass of pydantic.BaseModel"):
         create_str_enum_from_discriminated_type_union("TestEnum", type_union, "column_type")
 
 
-def test_create_str_enum_from_type_union_missing_discriminator_field() -> None:
+def test_create_str_enum_from_type_union_invalid_discriminator_field() -> None:
     type_union = Union[StubModelA, StubModelWithoutDiscriminator]
 
-    with pytest.raises(ValueError, match="column_type is not a field of"):
+    with pytest.raises(InvalidDiscriminatorFieldError, match="'column_type' is not a field of"):
         create_str_enum_from_discriminated_type_union("TestEnum", type_union, "column_type")
 
+    with pytest.raises(InvalidDiscriminatorFieldError, match="'name' must be a Literal type"):
+        create_str_enum_from_discriminated_type_union("TestEnum", type_union, "name")
+
 
 def test_create_str_enum_from_type_union_custom_discriminator_name() -> None:
     class StubModelE(BaseModel):