Generate a Python tagged union.
(
self,
union: Union,
indent: int = 0,
parent_stack: Optional[List[Message]] = None,
)
| 425 | return lines |
| 426 | |
| 427 | def generate_union( |
| 428 | self, |
| 429 | union: Union, |
| 430 | indent: int = 0, |
| 431 | parent_stack: Optional[List[Message]] = None, |
| 432 | ) -> List[str]: |
| 433 | """Generate a Python tagged union.""" |
| 434 | lines: List[str] = [] |
| 435 | ind = " " * indent |
| 436 | parent_path = "" |
| 437 | if parent_stack: |
| 438 | parent_path = ".".join([msg.name for msg in parent_stack]) |
| 439 | case_enum = f"{union.name}Case" |
| 440 | case_enum_ref = f"{parent_path}.{case_enum}" if parent_path else case_enum |
| 441 | union_ref = f"{parent_path}.{union.name}" if parent_path else union.name |
| 442 | |
| 443 | lines.append(f"{ind}class {case_enum}(Enum):") |
| 444 | for field in union.fields: |
| 445 | case_name = self.to_upper_snake_case(field.name) |
| 446 | lines.append(f"{ind} {case_name} = {field.number}") |
| 447 | lines.append("") |
| 448 | |
| 449 | comment = self.format_type_id_comment(union, f"{ind}#") |
| 450 | if comment: |
| 451 | lines.append(comment) |
| 452 | lines.append(f"{ind}class {union.name}(Union):") |
| 453 | lines.append(f'{ind} __slots__ = ("_case",)') |
| 454 | lines.append("") |
| 455 | lines.append( |
| 456 | f"{ind} def __init__(self, case: {case_enum_ref}, value: object) -> None:" |
| 457 | ) |
| 458 | lines.append(f"{ind} super().__init__(case.value, value)") |
| 459 | lines.append(f"{ind} self._case = case") |
| 460 | lines.append(f"{ind} self._validate()") |
| 461 | lines.append("") |
| 462 | |
| 463 | for field in union.fields: |
| 464 | method_name = self.safe_name(self.to_snake_case(field.name)) |
| 465 | case_name = self.to_upper_snake_case(field.name) |
| 466 | case_type = self.get_union_case_type(field, parent_stack) |
| 467 | lines.append(f"{ind} @classmethod") |
| 468 | lines.append( |
| 469 | f'{ind} def {method_name}(cls, v: {case_type}) -> "{union_ref}":' |
| 470 | ) |
| 471 | lines.append(f"{ind} return cls({case_enum_ref}.{case_name}, v)") |
| 472 | lines.append("") |
| 473 | |
| 474 | lines.append(f"{ind} @classmethod") |
| 475 | lines.append( |
| 476 | f'{ind} def _from_case_id(cls, case_id: int, value: object) -> "{union_ref}":' |
| 477 | ) |
| 478 | for field in union.fields: |
| 479 | case_name = self.to_upper_snake_case(field.name) |
| 480 | lines.append( |
| 481 | f"{ind} if case_id == {case_enum_ref}.{case_name}.value:" |
| 482 | ) |
| 483 | lines.append( |
| 484 | f"{ind} return cls({case_enum_ref}.{case_name}, value)" |
no test coverage detected