dynamically create the dd column type enum

Author: johnnygreco

src/data_designer/config/columns.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC
-from enum import Enum
 from typing import Literal, Optional, Type, Union
 
 from pydantic import BaseModel, Field, model_validator
@@ -15,21 +14,10 @@
 from .utils.code_lang import CodeLang
 from .utils.constants import REASONING_TRACE_COLUMN_POSTFIX
 from .utils.misc import assert_valid_jinja2_template, get_prompt_template_keywords
-from .utils.type_helpers import SAMPLER_PARAMS, resolve_string_enum
+from .utils.type_helpers import SAMPLER_PARAMS, create_str_enum_from_discriminated_type_union, resolve_string_enum
 from .validator_params import ValidatorParamsT, ValidatorType
 
 
-class DataDesignerColumnType(str, Enum):
-    SAMPLER = "sampler"
-    LLM_TEXT = "llm-text"
-    LLM_CODE = "llm-code"
-    LLM_STRUCTURED = "llm-structured"
-    LLM_JUDGE = "llm-judge"
-    EXPRESSION = "expression"
-    VALIDATION = "validation"
-    SEED_DATASET = "seed-dataset"
-
-
 class SingleColumnConfig(ConfigBase, ABC):
     name: str
     drop: bool = False
@@ -143,6 +131,25 @@ class SeedDatasetColumnConfig(SingleColumnConfig):
     column_type: Literal["seed-dataset"] = "seed-dataset"
 
 
+ColumnConfigT: TypeAlias = Union[
+    ExpressionColumnConfig,
+    LLMCodeColumnConfig,
+    LLMJudgeColumnConfig,
+    LLMStructuredColumnConfig,
+    LLMTextColumnConfig,
+    SamplerColumnConfig,
+    SeedDatasetColumnConfig,
+    ValidationColumnConfig,
+]
+
+
+DataDesignerColumnType = create_str_enum_from_discriminated_type_union(
+    enum_name="DataDesignerColumnType",
+    type_union=ColumnConfigT,
+    discriminator_field_name="column_type",
+)
+
+
 COLUMN_TYPE_EMOJI_MAP = {
     "general": "⚛️",  # possible analysis column type
     DataDesignerColumnType.EXPRESSION: "🧩",
@@ -156,18 +163,6 @@ class SeedDatasetColumnConfig(SingleColumnConfig):
 }
 
 
-ColumnConfigT: TypeAlias = Union[
-    ExpressionColumnConfig,
-    LLMCodeColumnConfig,
-    LLMJudgeColumnConfig,
-    LLMStructuredColumnConfig,
-    LLMTextColumnConfig,
-    SamplerColumnConfig,
-    SeedDatasetColumnConfig,
-    ValidationColumnConfig,
-]
-
-
 def column_type_is_in_dag(column_type: Union[str, DataDesignerColumnType]) -> bool:
     column_type = resolve_string_enum(column_type, DataDesignerColumnType)
     return column_type in [

src/data_designer/config/utils/type_helpers.py

@@ -3,14 +3,53 @@
 
 from enum import Enum
 import inspect
-from typing import Any, Type
+from typing import Any, Type, Union, get_args
 
 from pydantic import BaseModel
 
 from .. import sampler_params
 from .errors import InvalidEnumValueError
 
 
+class StrEnum(str, Enum):
+    pass
+
+
+def create_str_enum_from_discriminated_type_union(
+    enum_name: str,
+    type_union: Type[Union[BaseModel, ...]],
+    discriminator_field_name: str,
+) -> StrEnum:
+    """Create a string enum from a type union.
+
+    The type union is assumed to be a union of configs (Pydantic models) that have a discriminator field,
+    which must be a Literal string type - e.g., Literal["expression"].
+
+    Args:
+        enum_name: Name of the StrEnum.
+        type_union: Type union of configs (Pydantic models).
+        discriminator_field_name: Name of the discriminator field.
+
+    Returns:
+        StrEnum with values being the discriminator field values of the configs in the type union.
+
+    Example:
+        DataDesignerColumnType = create_str_enum_from_discriminated_type_union(
+            enum_name="DataDesignerColumnType",
+            type_union=ColumnConfigT,
+            discriminator_field_name="column_type",
+        )
+    """
+    discriminator_field_values = []
+    for model in type_union.__args__:
+        if not issubclass(model, BaseModel):
+            raise ValueError(f"🛑 {model} is not a Pydantic model.")
+        if discriminator_field_name not in model.model_fields:
+            raise ValueError(f"🛑 {discriminator_field_name} is not a field of {model}.")
+        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)})
+
+
 def get_sampler_params() -> dict[str, Type[BaseModel]]:
     """Returns a dictionary of sampler parameter classes."""
     params_cls_list = [

tests/config/utils/test_type_helpers.py

@@ -2,17 +2,127 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from enum import Enum
+from typing import Literal, Union
 
+from pydantic import BaseModel
 import pytest
 
 from data_designer.config.utils.errors import InvalidEnumValueError
-from data_designer.config.utils.type_helpers import SAMPLER_PARAMS, get_sampler_params, resolve_string_enum
+from data_designer.config.utils.type_helpers import (
+    SAMPLER_PARAMS,
+    create_str_enum_from_discriminated_type_union,
+    get_sampler_params,
+    resolve_string_enum,
+)
 
 
 class StubTestEnum(str, Enum):
     TEST = "test"
 
 
+class StubModelA(BaseModel):
+    column_type: Literal["type-a", "type-a-alt"] = "type-a"
+    name: str
+
+
+class StubModelB(BaseModel):
+    column_type: Literal["type-b"] = "type-b"
+    value: int
+
+
+class StubModelC(BaseModel):
+    column_type: Literal["type-c-with-dashes"] = "type-c-with-dashes"
+    data: str
+
+
+class StubModelWithoutDiscriminator(BaseModel):
+    name: str
+    value: int
+
+
+class NotAModel:
+    column_type: str = "not-a-model"
+
+
+def test_create_str_enum_from_type_union_basic() -> None:
+    type_union = Union[StubModelA, StubModelB]
+    result = create_str_enum_from_discriminated_type_union("TestEnum", type_union, "column_type")
+
+    assert issubclass(result, Enum)
+    assert issubclass(result, str)
+    assert hasattr(result, "TYPE_A")
+    assert hasattr(result, "TYPE_A_ALT")
+    assert hasattr(result, "TYPE_B")
+    assert result.TYPE_A.value == "type-a"
+    assert result.TYPE_A_ALT.value == "type-a-alt"
+    assert result.TYPE_B.value == "type-b"
+    assert len(result) == 3
+
+
+def test_create_str_enum_from_type_union_with_dashes() -> None:
+    type_union = Union[StubModelC, StubModelA]
+    result = create_str_enum_from_discriminated_type_union("TestEnum", type_union, "column_type")
+
+    assert hasattr(result, "TYPE_C_WITH_DASHES")
+    assert result.TYPE_C_WITH_DASHES.value == "type-c-with-dashes"
+
+
+def test_create_str_enum_from_type_union_multiple_models() -> None:
+    type_union = Union[StubModelA, StubModelB, StubModelC]
+    result = create_str_enum_from_discriminated_type_union("TestEnum", type_union, "column_type")
+
+    assert len(result) == 4
+    assert hasattr(result, "TYPE_A")
+    assert hasattr(result, "TYPE_A_ALT")
+    assert hasattr(result, "TYPE_B")
+    assert hasattr(result, "TYPE_C_WITH_DASHES")
+
+
+def test_create_str_enum_from_type_union_duplicate_values() -> None:
+    class StubModelD(BaseModel):
+        column_type: Literal["type-a"] = "type-a"
+        extra: str
+
+    type_union = Union[StubModelA, StubModelD]
+    result = create_str_enum_from_discriminated_type_union("TestEnum", type_union, "column_type")
+
+    assert len(result) == 2
+    assert hasattr(result, "TYPE_A")
+    assert hasattr(result, "TYPE_A_ALT")
+
+
+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"):
+        create_str_enum_from_discriminated_type_union("TestEnum", type_union, "column_type")
+
+
+def test_create_str_enum_from_type_union_missing_discriminator_field() -> None:
+    type_union = Union[StubModelA, StubModelWithoutDiscriminator]
+
+    with pytest.raises(ValueError, match="column_type is not a field of"):
+        create_str_enum_from_discriminated_type_union("TestEnum", type_union, "column_type")
+
+
+def test_create_str_enum_from_type_union_custom_discriminator_name() -> None:
+    class StubModelE(BaseModel):
+        type_field: Literal["custom-type"] = "custom-type"
+        name: str
+
+    class StubModelF(BaseModel):
+        type_field: Literal["another-type"] = "another-type"
+        value: int
+
+    type_union = Union[StubModelE, StubModelF]
+    result = create_str_enum_from_discriminated_type_union("TestEnum", type_union, "type_field")
+
+    assert hasattr(result, "CUSTOM_TYPE")
+    assert result.CUSTOM_TYPE.value == "custom-type"
+    assert hasattr(result, "ANOTHER_TYPE")
+    assert result.ANOTHER_TYPE.value == "another-type"
+
+
 def test_get_sampler_params():
     expected_sampler_keys = {
         "bernoulli",