commit 0de6577b7d5702a133f0eade4ddf6b94893e975a
author Oliver Bristow <evilumbrella+github@gmail.com> Thu Feb 28 00:39:51 2019 +0000
committer Jie Luo <anandolee@gmail.com> Wed Feb 27 16:39:51 2019 -0800
tree 68057b6b42fb381f22c044cae9d32e13c200fc7a
parent 05e81cb8f41db6e124ae8bcd52cc9abff530a31a

Add EnumTypeWrapper.__getattr__ to access values (#5234)

python/google/protobuf/internal/enum_type_wrapper.py

diff --git a/python/google/protobuf/internal/enum_type_wrapper.py b/python/google/protobuf/internal/enum_type_wrapper.py
index 77e4be0..cc9e95d 100644
--- a/python/google/protobuf/internal/enum_type_wrapper.py
+++ b/python/google/protobuf/internal/enum_type_wrapper.py
@@ -87,3 +87,9 @@
     """
     return [(value_descriptor.name, value_descriptor.number)
             for value_descriptor in self._enum_type.values]
+
+  def __getattr__(self, name):
+    """Returns the value coresponding to the given enum name."""
+    if name in self._enum_type.values_by_name:
+      return self._enum_type.values_by_name[name].number
+    raise AttributeError

python/google/protobuf/internal/reflection_test.py

diff --git a/python/google/protobuf/internal/reflection_test.py b/python/google/protobuf/internal/reflection_test.py
index 90d2fe3..b87eaba 100755
--- a/python/google/protobuf/internal/reflection_test.py
+++ b/python/google/protobuf/internal/reflection_test.py
@@ -804,30 +804,64 @@
   def testEnum_Value(self):
     self.assertEqual(unittest_pb2.FOREIGN_FOO,
                      unittest_pb2.ForeignEnum.Value('FOREIGN_FOO'))
+    self.assertEqual(unittest_pb2.FOREIGN_FOO,
+                     unittest_pb2.ForeignEnum.FOREIGN_FOO)
+
     self.assertEqual(unittest_pb2.FOREIGN_BAR,
                      unittest_pb2.ForeignEnum.Value('FOREIGN_BAR'))
+    self.assertEqual(unittest_pb2.FOREIGN_BAR,
+                     unittest_pb2.ForeignEnum.FOREIGN_BAR)
+
     self.assertEqual(unittest_pb2.FOREIGN_BAZ,
                      unittest_pb2.ForeignEnum.Value('FOREIGN_BAZ'))
+    self.assertEqual(unittest_pb2.FOREIGN_BAZ,
+                     unittest_pb2.ForeignEnum.FOREIGN_BAZ)
+
     self.assertRaises(ValueError,
                       unittest_pb2.ForeignEnum.Value, 'FO')
+    with self.assertRaises(AttributeError):
+      unittest_pb2.ForeignEnum.FO
 
     proto = unittest_pb2.TestAllTypes()
     self.assertEqual(proto.FOO,
                      proto.NestedEnum.Value('FOO'))
     self.assertEqual(proto.FOO,
+                     proto.NestedEnum.FOO)
+
+    self.assertEqual(proto.FOO,
                      unittest_pb2.TestAllTypes.NestedEnum.Value('FOO'))
+    self.assertEqual(proto.FOO,
+                     unittest_pb2.TestAllTypes.NestedEnum.FOO)
+
     self.assertEqual(proto.BAR,
                      proto.NestedEnum.Value('BAR'))
     self.assertEqual(proto.BAR,
+                     proto.NestedEnum.BAR)
+
+    self.assertEqual(proto.BAR,
                      unittest_pb2.TestAllTypes.NestedEnum.Value('BAR'))
+    self.assertEqual(proto.BAR,
+                     unittest_pb2.TestAllTypes.NestedEnum.BAR)
+
     self.assertEqual(proto.BAZ,
                      proto.NestedEnum.Value('BAZ'))
     self.assertEqual(proto.BAZ,
+                     proto.NestedEnum.BAZ)
+
+    self.assertEqual(proto.BAZ,
                      unittest_pb2.TestAllTypes.NestedEnum.Value('BAZ'))
+    self.assertEqual(proto.BAZ,
+                     unittest_pb2.TestAllTypes.NestedEnum.BAZ)
+
     self.assertRaises(ValueError,
                       proto.NestedEnum.Value, 'Foo')
+    with self.assertRaises(AttributeError):
+      proto.NestedEnum.Value.Foo
+
     self.assertRaises(ValueError,
                       unittest_pb2.TestAllTypes.NestedEnum.Value, 'Foo')
+    with self.assertRaises(AttributeError):
+      unittest_pb2.TestAllTypes.NestedEnum.Value.Foo
 
   def testEnum_KeysAndValues(self):
     self.assertEqual(['FOREIGN_FOO', 'FOREIGN_BAR', 'FOREIGN_BAZ'],