fix(types): generate schema enums as StrEnum (#915)

Author: bokelley

scripts/post_generate_fixes.py

@@ -6,12 +6,13 @@
 This script applies necessary modifications to generated files that cannot be
 handled by datamodel-code-generator directly:
 
-1. Adds model_validators to types requiring mutual exclusivity checks
-2. Fixes self-referential RootModel type annotations
-3. Fixes BrandManifest forward references
-4. Adds deprecated=True to fields marked deprecated in JSON schema
-5. Unwraps specified RootModel unions to plain Union type aliases (#155)
-6. Widens canceled: Literal[True] = True on request types to | None = None (#641)
+1. Rewrites generated string enums to StrEnum
+2. Adds model_validators to types requiring mutual exclusivity checks
+3. Fixes self-referential RootModel type annotations
+4. Fixes BrandManifest forward references
+5. Adds deprecated=True to fields marked deprecated in JSON schema
+6. Unwraps specified RootModel unions to plain Union type aliases (#155)
+7. Widens canceled: Literal[True] = True on request types to | None = None (#641)
 """
 
 from __future__ import annotations
@@ -48,6 +49,8 @@ def _load_resolve_bundle_key():
 
 _PROTOCOL_ENVELOPE_IMPORT = "from ..core.protocol_envelope import ProtocolEnvelope\n"
 _VERSION_ENVELOPE_IMPORT = "from ..core.version_envelope import AdcpVersionEnvelope\n"
+_STR_ENUM_MEMBER_ASSIGNMENT_IGNORE = "  # type: ignore[assignment]"
+_STR_ATTRIBUTE_NAMES = set(dir(str))
 
 
 def _sync_protocol_envelope_import(source: str) -> str:
@@ -82,6 +85,90 @@ def add_model_validator_to_product():
     print("  product.py validation: no fixes needed (Pydantic handles discriminated unions)")
 
 
+def _ignore_strenum_member_method_collisions(source: str) -> tuple[str, int]:
+    """Suppress mypy for StrEnum members that intentionally shadow str methods."""
+    lines = source.splitlines()
+    updated: list[str] = []
+    class_indent: int | None = None
+    ignores_added = 0
+
+    for line in lines:
+        stripped = line.lstrip()
+        indent = len(line) - len(stripped)
+
+        class_match = re.match(r"class\s+\w+\(StrEnum\):", stripped)
+        if class_match is not None:
+            class_indent = indent
+            updated.append(line)
+            continue
+
+        if class_indent is not None and stripped and indent <= class_indent:
+            class_indent = None
+
+        if class_indent is not None and indent == class_indent + 4:
+            assignment_match = re.match(r"([A-Za-z_]\w*)\s*=", stripped)
+            if (
+                assignment_match is not None
+                and assignment_match.group(1) in _STR_ATTRIBUTE_NAMES
+                and _STR_ENUM_MEMBER_ASSIGNMENT_IGNORE not in line
+            ):
+                line = f"{line}{_STR_ENUM_MEMBER_ASSIGNMENT_IGNORE}"
+                ignores_added += 1
+
+        updated.append(line)
+
+    return "\n".join(updated) + ("\n" if source.endswith("\n") else ""), ignores_added
+
+
+def rewrite_generated_enums_to_strenum() -> None:
+    """Make all generated schema enums inherit from StrEnum.
+
+    datamodel-code-generator emits plain ``Enum`` classes for string-valued
+    JSON Schema enums. The generated enum members should behave like their wire
+    values for equality, hashing, formatting, and ``str()`` without widening or
+    narrowing any model field annotations.
+    """
+    files_changed = 0
+    classes_changed = 0
+    ignores_added = 0
+
+    for path in OUTPUT_DIR.rglob("*.py"):
+        source = path.read_text()
+        if (
+            "(Enum):" not in source
+            and "from enum import Enum" not in source
+            and "(StrEnum):" not in source
+        ):
+            continue
+
+        updated = source.replace(
+            "from enum import Enum, IntEnum\n",
+            "from enum import IntEnum\nfrom adcp.types._str_enum import StrEnum\n",
+        )
+        updated = updated.replace(
+            "from enum import IntEnum, Enum\n",
+            "from enum import IntEnum\nfrom adcp.types._str_enum import StrEnum\n",
+        )
+        updated = updated.replace(
+            "from enum import Enum\n",
+            "from adcp.types._str_enum import StrEnum\n",
+        )
+        updated, changed = re.subn(r"\((?:str,\s*)?Enum\):", "(StrEnum):", updated)
+        updated, file_ignores_added = _ignore_strenum_member_method_collisions(updated)
+
+        if updated != source:
+            path.write_text(updated)
+            files_changed += 1
+            classes_changed += changed
+            ignores_added += file_ignores_added
+
+    print(
+        f"  generated enums rewritten to StrEnum "
+        f"({classes_changed} classes across {files_changed} files, "
+        f"{ignores_added} member type ignore(s))"
+    )
+
+
 def fix_preview_render_self_reference():
     """Fix self-referential RootModel in preview_render.py."""
     preview_file = OUTPUT_DIR / "creative" / "preview_render.py"
@@ -3297,6 +3384,7 @@ def main():
         fix_response_payload_jws_required_literals,
         fix_verify_brand_claim_models,
         fix_signal_coverage_forecast_point_types,
+        rewrite_generated_enums_to_strenum,
         strip_extra_blank_lines_at_eof,
     ]
     for fix in fixes:

src/adcp/types/_str_enum.py

@@ -0,0 +1,19 @@
+"""Python-version-compatible StrEnum export for generated schema enums."""
+
+from __future__ import annotations
+
+import sys
+from enum import Enum
+
+if sys.version_info >= (3, 11):
+    from enum import StrEnum as StrEnum
+else:
+
+    class StrEnum(str, Enum):
+        """Backport the stdlib StrEnum behavior needed by generated enums."""
+
+        def __str__(self) -> str:
+            return str.__str__(self)
+
+        def __format__(self, format_spec: str) -> str:
+            return str.__format__(self, format_spec)

src/adcp/types/generated_poc/a2ui/si_catalog.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from enum import Enum
+from adcp.types._str_enum import StrEnum
 from typing import Annotated, Literal
 
 from adcp.types.base import AdCPBaseModel
@@ -13,7 +13,7 @@
 from . import bound_value
 
 
-class Component(Enum):
+class Component(StrEnum):
     Text = 'Text'
     Button = 'Button'
     Link = 'Link'
@@ -36,7 +36,7 @@ class SiComponentCatalog(AdCPBaseModel):
     ] = None
 
 
-class Variant(Enum):
+class Variant(StrEnum):
     body = 'body'
     heading = 'heading'
     caption = 'caption'
@@ -56,7 +56,7 @@ class Action(AdCPBaseModel):
     ] = None
 
 
-class Variant5(Enum):
+class Variant5(StrEnum):
     primary = 'primary'
     secondary = 'secondary'
     text = 'text'
@@ -122,7 +122,7 @@ class Template(AdCPBaseModel):
     componentId: Annotated[str, Field(description='ID of component to use as template')]
 
 
-class Layout(Enum):
+class Layout(StrEnum):
     vertical = 'vertical'
     horizontal = 'horizontal'
     grid = 'grid'
@@ -137,16 +137,16 @@ class List(AdCPBaseModel):
     layout: Layout | None = Layout.vertical
 
 
-class Align(Enum):
+class Align(StrEnum):
     start = 'start'
-    center = 'center'
+    center = 'center'  # type: ignore[assignment]
     end = 'end'
     stretch = 'stretch'
 
 
-class Justify(Enum):
+class Justify(StrEnum):
     start = 'start'
-    center = 'center'
+    center = 'center'  # type: ignore[assignment]
     end = 'end'
     between = 'between'
     around = 'around'
@@ -165,7 +165,7 @@ class Column(AdCPBaseModel):
     align: Align | None = Align.stretch
 
 
-class Type(Enum):
+class Type(StrEnum):
     mcp = 'mcp'
     a2a = 'a2a'
 

src/adcp/types/generated_poc/aao/agent_publishers.py

@@ -4,21 +4,21 @@
 
 from __future__ import annotations
 
-from enum import Enum
+from adcp.types._str_enum import StrEnum
 from typing import Annotated
 
 from adcp.types.base import AdCPBaseModel
 from pydantic import AnyUrl, AwareDatetime, ConfigDict, Field
 
 
-class DiscoveryMethod(Enum):
+class DiscoveryMethod(StrEnum):
     direct = 'direct'
     authoritative_location = 'authoritative_location'
     adagents_authoritative = 'adagents_authoritative'
     ads_txt_managerdomain = 'ads_txt_managerdomain'
 
 
-class Status(Enum):
+class Status(StrEnum):
     authorized = 'authorized'
     revoked = 'revoked'
 

src/adcp/types/generated_poc/account/list_accounts_request.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from enum import Enum
+from adcp.types._str_enum import StrEnum
 from typing import Annotated
 
 from pydantic import ConfigDict, Field
@@ -16,7 +16,7 @@
 from ..core.version_envelope import AdcpVersionEnvelope
 
 
-class Status(Enum):
+class Status(StrEnum):
     active = 'active'
     pending_approval = 'pending_approval'
     rejected = 'rejected'

src/adcp/types/generated_poc/adagents.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from enum import Enum
+from adcp.types._str_enum import StrEnum
 from typing import Any, Annotated, Literal
 
 from adcp.types.base import AdCPBaseModel
@@ -94,7 +94,7 @@ class Contact(AdCPBaseModel):
     ] = None
 
 
-class Reason(Enum):
+class Reason(StrEnum):
     relationship_ended = 'relationship_ended'
     compliance_violation = 'compliance_violation'
     publisher_request = 'publisher_request'
@@ -144,7 +144,7 @@ class PlacementTags(AdCPBaseModel):
     ]
 
 
-class DelegationType(Enum):
+class DelegationType(StrEnum):
     direct = 'direct'
     delegated = 'delegated'
     ad_network = 'ad_network'