feat Add Enum datatype

Author: emekaokoli19

guardrails/datatypes.py

@@ -6,6 +6,7 @@
 from typing import Any, Dict, Iterable
 from typing import List as TypedList
 from typing import Optional, Sequence, Type, TypeVar, Union
+from enum import Enum
 
 from dateutil.parser import parse
 from lxml import etree as ET
@@ -386,6 +387,50 @@ class Percentage(ScalarType):
     tag = "percentage"
 
 
+@register_type("enum")
+class Enum(ScalarType):
+    """Element tag: `<enum>`"""
+
+    tag = "enum"
+
+    def __init__(
+        self,
+        children: Dict[str, Any],
+        validators_attr: ValidatorsAttr,
+        optional: bool,
+        name: Optional[str],
+        description: Optional[str],
+        enum_values: TypedList[str],
+    ) -> None:
+        super().__init__(children, validators_attr, optional, name, description)
+        self.enum_values = enum_values
+
+    def from_str(self, s: str) -> Optional[str]:
+        """Create an Enum from a string."""
+        if s is None:
+            return None
+        if s not in self.enum_values:
+            raise ValueError(f"Invalid enum value: {s}")
+        return s
+
+    @classmethod
+    def from_xml(
+        cls,
+        enum_values: TypedList[str],
+        validators: Sequence[ValidatorSpec],
+        description: Optional[str] = None,
+        strict: bool = False,
+    ) -> "Enum":
+        return cls(
+            children={},
+            validators_attr=ValidatorsAttr.from_validators(validators, cls.tag, strict),
+            optional=False,
+            name=None,
+            description=description,
+            enum_values=enum_values,
+        )
+
+
 @register_type("list")
 class List(NonScalarType):
     """Element tag: `<list>`"""

guardrails/utils/pydantic_utils/v1.py

@@ -3,6 +3,7 @@
 import warnings
 from copy import deepcopy
 from datetime import date, time
+from enum import Enum
 from typing import Any, Callable, Dict, Optional, Type, Union, get_args, get_origin
 
 from pydantic import BaseModel, validator
@@ -20,6 +21,7 @@
 from guardrails.datatypes import Object as ObjectDataType
 from guardrails.datatypes import String as StringDataType
 from guardrails.datatypes import Time as TimeDataType
+from guardrails.datatypes import Enum as EnumDataType
 from guardrails.validator_base import Validator
 from guardrails.validatorsattr import ValidatorsAttr
 
@@ -66,6 +68,16 @@ def is_dict(type_annotation: Any) -> bool:
         return True
     return False
 
+def is_enum(type_annotation: Any) -> bool:
+    """Check if a type_annotation is an enum"""
+
+    type_annotation = prepare_type_annotation(type_annotation)
+
+    if is_pydantic_base_model(type_annotation):
+        return False
+    if issubclass(type_annotation, Enum):
+        return True 
+    return False
 
 def prepare_type_annotation(type_annotation: Union[ModelField, Type]) -> Type:
     """Get the raw type annotation that can be used for downstream processing.
@@ -262,6 +274,8 @@ def field_to_datatype(field: Union[ModelField, Type]) -> Type[DataType]:
         return ListDataType
     elif is_dict(type_annotation):
         return ObjectDataType
+    elif is_enum(type_annotation):
+        return EnumDataType
     elif type_annotation == bool:
         return BooleanDataType
     elif type_annotation == date:
@@ -301,6 +315,9 @@ def convert_pydantic_model_to_datatype(
 
     model_fields = add_pydantic_validators_as_guardrails_validators(model)
 
+    # Use inline import to avoid circular dependency
+    from guardrails.validators import ValidChoices
+
     children = {}
     for field_name, field in model_fields.items():
         if field_name in excluded_fields:
@@ -356,6 +373,15 @@ def convert_pydantic_model_to_datatype(
                 strict=strict,
                 discriminator_key=discriminator,
             )
+        elif target_datatype == EnumDataType:
+            valid_choices = type_annotation._member_names_
+            field.field_info.extra["validators"] = [ValidChoices(choices=valid_choices)]
+            return pydantic_field_to_datatype(
+                EnumDataType,
+                field,
+                strict=strict,
+                enum_values=valid_choices
+            )
         elif isinstance(field.type_, type) and issubclass(field.type_, BaseModel):
             children[field_name] = convert_pydantic_model_to_datatype(
                 field, datatype=target_datatype, strict=strict

guardrails/utils/pydantic_utils/v2.py

@@ -2,8 +2,8 @@
 import warnings
 from copy import deepcopy
 from datetime import date, time
+from enum import Enum
 from typing import Any, Callable, Dict, List, Optional, Type, TypeVar, Union, get_args
-
 from pydantic import BaseModel, ConfigDict, HttpUrl, field_validator
 from pydantic.fields import FieldInfo
 
@@ -21,6 +21,7 @@
 from guardrails.datatypes import PythonCode as PythonCodeDataType
 from guardrails.datatypes import String as StringDataType
 from guardrails.datatypes import Time as TimeDataType
+from guardrails.datatypes import Enum as EnumDataType
 from guardrails.validator_base import Validator
 from guardrails.validatorsattr import ValidatorsAttr
 
@@ -87,6 +88,16 @@ def is_dict(type_annotation: Any) -> bool:
         return True
     return False
 
+def is_enum(type_annotation: Any) -> bool:
+    """Check if a type_annotation is an enum"""
+
+    type_annotation = prepare_type_annotation(type_annotation)
+
+    if is_pydantic_base_model(type_annotation):
+        return False
+    if issubclass(type_annotation, Enum):
+        return True 
+    return False
 
 def _create_bare_model(model: Type[BaseModel]) -> Type[BaseModel]:
     class BareModel(BaseModel):
@@ -277,6 +288,8 @@ def field_to_datatype(field: Union[FieldInfo, Type]) -> Type[DataType]:
         return ListDataType
     elif is_dict(type_annotation):
         return ObjectDataType
+    elif is_enum(type_annotation):
+        return EnumDataType
     elif type_annotation == bool:
         return BooleanDataType
     elif type_annotation == date:
@@ -322,6 +335,9 @@ def convert_pydantic_model_to_datatype(
 
     model_fields = add_pydantic_validators_as_guardrails_validators(model)
 
+    # Use inline import to avoid circular dependency
+    from guardrails.validators import ValidChoices
+
     children = {}
     for field_name, field in model_fields.items():
         if field_name in excluded_fields:
@@ -382,6 +398,15 @@ def convert_pydantic_model_to_datatype(
                 discriminator_key=discriminator,
                 name=field_name,
             )
+        elif target_datatype == EnumDataType:
+            valid_choices = type_annotation._member_names_
+            field.json_schema_extra["validators"] = [ValidChoices(choices=valid_choices)]
+            return pydantic_field_to_datatype(
+                EnumDataType,
+                field,
+                strict=strict,
+                enum_values=valid_choices
+            )
         elif is_pydantic_base_model(field.annotation):
             children[field_name] = convert_pydantic_model_to_datatype(
                 field,