Map enum int's into Enums redux (#293)

Re-implement Enum to be faster along with being an open set

---------
Co-authored-by: ydylla <ydylla@gmail.com>
This commit is contained in:
James Hilton-Balfe
2023-10-16 03:32:30 +01:00
committed by GitHub
parent 8659c51123
commit c82816b8be
6 changed files with 725 additions and 511 deletions

View File

@@ -1,5 +1,5 @@
import dataclasses
import enum
import enum as builtin_enum
import json
import math
import struct
@@ -45,7 +45,8 @@ from .casing import (
safe_snake_case,
snake_case,
)
from .grpc.grpclib_client import ServiceStub
from .enum import Enum as Enum
from .grpc.grpclib_client import ServiceStub as ServiceStub
if TYPE_CHECKING:
@@ -140,7 +141,7 @@ NEG_INFINITY = "-Infinity"
NAN = "NaN"
class Casing(enum.Enum):
class Casing(builtin_enum.Enum):
"""Casing constants for serialization."""
CAMEL = camel_case #: A camelCase sterilization function.
@@ -309,32 +310,6 @@ def map_field(
)
class Enum(enum.IntEnum):
"""
The base class for protobuf enumerations, all generated enumerations will inherit
from this. Bases :class:`enum.IntEnum`.
"""
@classmethod
def from_string(cls, name: str) -> "Enum":
"""Return the value which corresponds to the string name.
Parameters
-----------
name: :class:`str`
The name of the enum member to get
Raises
-------
:exc:`ValueError`
The member was not found in the Enum.
"""
try:
return cls._member_map_[name] # type: ignore
except KeyError as e:
raise ValueError(f"Unknown value {name} for enum {cls.__name__}") from e
def _pack_fmt(proto_type: str) -> str:
"""Returns a little-endian format string for reading/writing binary."""
return {
@@ -1168,7 +1143,7 @@ class Message(ABC):
return t
elif issubclass(t, Enum):
# Enums always default to zero.
return int
return t.try_value
elif t is datetime:
# Offsets are relative to 1970-01-01T00:00:00Z
return datetime_default_gen
@@ -1193,6 +1168,9 @@ class Message(ABC):
elif meta.proto_type == TYPE_BOOL:
# Booleans use a varint encoding, so convert it to true/false.
value = value > 0
elif meta.proto_type == TYPE_ENUM:
# Convert enum ints to python enum instances
value = self._betterproto.cls_by_field[field_name].try_value(value)
elif wire_type in (WIRE_FIXED_32, WIRE_FIXED_64):
fmt = _pack_fmt(meta.proto_type)
value = struct.unpack(fmt, value)[0]