소스 검색

Add EnumTypeWrapper.__getattr__ to access values (#5234)

Oliver Bristow 6 년 전
부모
커밋
0de6577b7d
2개의 변경된 파일40개의 추가작업 그리고 0개의 파일을 삭제
  1. 6 0
      python/google/protobuf/internal/enum_type_wrapper.py
  2. 34 0
      python/google/protobuf/internal/reflection_test.py

+ 6 - 0
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

+ 34 - 0
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'],