_enum.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  1. # Copyright (c) 2025-2026 Buf Technologies, Inc.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import annotations
  15. from enum import IntEnum
  16. from typing import TYPE_CHECKING
  17. if TYPE_CHECKING:
  18. from ._descriptors import DescEnum
  19. class Enum(IntEnum):
  20. """Base class for protobuf enumeration types.
  21. Protobuf enumerations are integer-based. This class inherits from `IntEnum`
  22. so enum values work as integers while maintaining a link to the protobuf
  23. descriptor.
  24. Generated enum classes inherit from this base class and define their enum
  25. values as class attributes.
  26. Examples:
  27. ```python
  28. from protobuf import Enum
  29. # Generated from:
  30. # enum Color {
  31. # COLOR_UNSPECIFIED = 0;
  32. # COLOR_RED = 1;
  33. # COLOR_GREEN = 2;
  34. # }
  35. class Color(Enum):
  36. UNSPECIFIED = 0 # Python name (prefix stripped)
  37. RED = 1
  38. GREEN = 2
  39. # Use like a regular int
  40. color = Color.RED
  41. print(color) # Prints: RED (Python member name)
  42. print(repr(color)) # Color.RED (Python qualified name)
  43. print(int(color)) # 1
  44. print(color + 1) # 2
  45. print(color == 1) # True
  46. # Access the descriptor
  47. desc = Color.desc()
  48. print(desc.name) # "Color"
  49. ```
  50. """
  51. if TYPE_CHECKING:
  52. _desc: DescEnum
  53. @classmethod
  54. def _missing_(cls, value: object) -> Enum:
  55. """Handle unknown enum values.
  56. Called by the enum metaclass for undefined values. Open enums must
  57. support unknown values, so this creates a pseudo-member.
  58. Args:
  59. value: The integer value to create an enum instance for.
  60. Returns:
  61. A new enum instance with the given value.
  62. Raises:
  63. TypeError: If the value is not an integer.
  64. """
  65. if not isinstance(value, int):
  66. msg = f"value must be an int, not {type(value).__name__}"
  67. raise TypeError(msg)
  68. if not cls._desc.open:
  69. msg = f"`{value}` is not a valid {cls.__name__}"
  70. raise ValueError(msg)
  71. # Create a pseudo-member for unknown values
  72. pseudo_member = int.__new__(cls, value)
  73. pseudo_member._name_ = None # type: ignore[attr-defined]
  74. pseudo_member._value_ = value # type: ignore[attr-defined]
  75. return pseudo_member # type: ignore[return-value]
  76. def __str__(self) -> str:
  77. """Return a string representation of the enum value.
  78. Returns the Python member name (e.g., "RED") for known values,
  79. or the integer as a string for unknown values.
  80. Returns:
  81. The Python member name if known, otherwise the integer as a string.
  82. Examples:
  83. ```python
  84. str(Color.RED) # 'RED'
  85. str(Color(99)) # '99'
  86. ```
  87. """
  88. if self._name_ is not None:
  89. return self._name_
  90. return str(int(self))
  91. def __repr__(self) -> str:
  92. """Return a detailed string representation of the enum value.
  93. Returns the qualified Python name (e.g., "Color.RED") for known values,
  94. or "Color(99)" format for unknown values.
  95. Returns:
  96. The qualified Python name if known, otherwise the class name with integer.
  97. Examples:
  98. ```python
  99. repr(Color.RED) # 'Color.RED'
  100. repr(Color(99)) # 'Color(99)'
  101. ```
  102. """
  103. if self._name_ is not None:
  104. return f"{self.__class__.__qualname__}.{self._name_}"
  105. return f"{self.__class__.__qualname__}({int(self)})"
  106. @classmethod
  107. def desc(cls) -> DescEnum:
  108. """Get the descriptor for this enumeration type.
  109. Returns:
  110. The DescEnum descriptor for this enum.
  111. Examples:
  112. ```python
  113. desc = Color.desc()
  114. desc.type_name # 'my.package.Color'
  115. desc.values[0].name # 'COLOR_UNSPECIFIED'
  116. ```
  117. """
  118. return cls._desc
  119. def enum_is_unknown(value: Enum, /) -> bool:
  120. """Check if an enum value is unknown (i.e., not defined in the protobuf schema)."""
  121. return value.name is None