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