Add t property to SumType

CQCL · May 9, 2024 · 1d09946 · 1d09946
1 parent f3d8bcc
commit 1d09946
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/hugr-py/src/hugr/serialization/tys.py b/hugr-py/src/hugr/serialization/tys.py
@@ -190,7 +190,12 @@ class GeneralSum(ConfiguredBaseModel):
 
 
 class SumType(RootModel):
-    root: Union[UnitSum, GeneralSum] = Field(discriminator="s")
+    root: Annotated[Union[UnitSum, GeneralSum], Field(discriminator="s")]
+
+    # This seems to be required for nested discriminated unions to work
+    @property
+    def t(self) -> str:
+        return self.root.t
 
     class Config:
         json_schema_extra = {"required": ["s"]}
@@ -318,7 +323,8 @@ class Type(RootModel):
     root: Annotated[
         Qubit | Variable | USize | FunctionType | Array | SumType | Opaque | Alias,
         WrapValidator(_json_custom_error_validator),
-    ] = Field(discriminator="t")
+        Field(discriminator="t"),
+    ]
 
     class Config:
         json_schema_extra = {"required": ["t"]}