-
Notifications
You must be signed in to change notification settings - Fork 16.3k
Expand file tree
/
Copy pathenum_type_wrapper.py
More file actions
128 lines (108 loc) · 3.8 KB
/
Copy pathenum_type_wrapper.py
File metadata and controls
128 lines (108 loc) · 3.8 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
# Protocol Buffers - Google's data interchange format
# Copyright 2008 Google Inc. All rights reserved.
#
# Use of this source code is governed by a BSD-style
# license that can be found in the LICENSE file or at
# https://developers.google.com/open-source/licenses/bsd
"""A simple wrapper around enum types to expose utility functions.
Instances are created as properties with the same name as the enum they wrap
on proto classes. For usage, see:
reflection_test.py
"""
import sys
__author__ = 'rabsatt@google.com (Kevin Rabsatt)'
class EnumTypeWrapper(object):
"""A utility for finding the names of enum values."""
DESCRIPTOR = None
# This is a type alias, which mypy typing stubs can type as
# a genericized parameter constrained to an int, allowing subclasses
# to be typed with more constraint in .pyi stubs
# Eg.
# def MyGeneratedEnum(Message):
# ValueType = NewType('ValueType', int)
# def Name(self, number: MyGeneratedEnum.ValueType) -> str
ValueType = int
def __init__(self, enum_type):
"""Inits EnumTypeWrapper with an EnumDescriptor."""
self._enum_type = enum_type
self.DESCRIPTOR = enum_type # pylint: disable=invalid-name
def Name(self, number): # pylint: disable=invalid-name
"""Returns a string containing the name of an enum value."""
try:
return self._enum_type.values_by_number[number].name
except KeyError:
pass # fall out to break exception chaining
if not isinstance(number, int):
raise TypeError(
'Enum value for {} must be an int, but got {} {!r}.'.format(
self._enum_type.name, type(number), number
)
)
else:
# repr here to handle the odd case when you pass in a boolean.
raise ValueError(
'Enum {} has no name defined for value {!r}'.format(
self._enum_type.name, number
)
)
def Value(self, name): # pylint: disable=invalid-name
"""Returns the value corresponding to the given enum name."""
try:
return self._enum_type.values_by_name[name].number
except KeyError:
pass # fall out to break exception chaining
raise ValueError(
'Enum {} has no value defined for name {!r}'.format(
self._enum_type.name, name
)
)
def keys(self):
"""Return a list of the string names in the enum.
Returns:
A list of strs, in the order they were defined in the .proto file.
"""
return [
value_descriptor.name for value_descriptor in self._enum_type.values
]
def values(self):
"""Return a list of the integer values in the enum.
Returns:
A list of ints, in the order they were defined in the .proto file.
"""
return [
value_descriptor.number for value_descriptor in self._enum_type.values
]
def items(self):
"""Return a list of the (name, value) pairs of the enum.
Returns:
A list of (str, int) pairs, in the order they were defined
in the .proto file.
"""
return [
(value_descriptor.name, value_descriptor.number)
for value_descriptor in self._enum_type.values
]
def __getattr__(self, name):
"""Returns the value corresponding to the given enum name."""
try:
return (
super(EnumTypeWrapper, self)
.__getattribute__('_enum_type')
.values_by_name[name]
.number
)
except KeyError:
pass # fall out to break exception chaining
raise AttributeError(
'Enum {} has no value defined for name {!r}'.format(
self._enum_type.name, name
)
)
def __or__(self, other):
"""Returns the union type of self and other."""
if sys.version_info >= (3, 10):
return type(self) | other
else:
raise NotImplementedError(
'You may not use | on EnumTypes (or classes) below python 3.10'
)