From 0de6577b7d5702a133f0eade4ddf6b94893e975a Mon Sep 17 00:00:00 2001 From: Oliver Bristow Date: Thu, 28 Feb 2019 00:39:51 +0000 Subject: [PATCH] Add EnumTypeWrapper.__getattr__ to access values (#5234) --- .../protobuf/internal/enum_type_wrapper.py | 6 ++++ .../protobuf/internal/reflection_test.py | 34 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/python/google/protobuf/internal/enum_type_wrapper.py b/python/google/protobuf/internal/enum_type_wrapper.py index 77e4be01a3..cc9e95daa0 100644 --- a/python/google/protobuf/internal/enum_type_wrapper.py +++ b/python/google/protobuf/internal/enum_type_wrapper.py @@ -87,3 +87,9 @@ class EnumTypeWrapper(object): """ 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 diff --git a/python/google/protobuf/internal/reflection_test.py b/python/google/protobuf/internal/reflection_test.py index 90d2fe3cc6..b87eabaa9b 100755 --- a/python/google/protobuf/internal/reflection_test.py +++ b/python/google/protobuf/internal/reflection_test.py @@ -804,30 +804,64 @@ class ReflectionTest(BaseTestCase): 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'],