Merge pull request #262 from jayvdb/use-enum

_StructModuleWhich: Use enum

capnp/lib/capnp.pyx

@@ -22,6 +22,7 @@ from libc.string cimport memcpy
 import array
 import asyncio
 import collections as _collections
+import enum as _enum
 import inspect as _inspect
 import os as _os
 import random as _random
@@ -1392,6 +1393,8 @@ cdef class _DynamicStructBuilder:
 
         :Raises: :exc:`KjException` if the field isn't in this struct
         """
+        if isinstance(field, _StructModuleWhich):
+            field = field.name[0].lower() + field.name[1:]
         if size is None:
             return to_python_builder(self.thisptr.init(field), self._parent)
         else:
@@ -3152,8 +3155,12 @@ cdef _new_message(self, kwargs, num_first_segment_words):
     return msg
 
 
-class _StructModuleWhich(object):
+class _StructModuleWhich(_enum.Enum):
-    pass
+    def __eq__(self, other):
+        if isinstance(other, int):
+            return self.value == other
+        else:
+            return self.name == other
 
 
 class _StructModule(object):
@@ -3170,17 +3177,19 @@ class _StructModule(object):
                 if field_schema.discriminantCount == 0:
                     sub_module = _StructModule(raw_schema, name)
                 else:
-                    sub_module = _StructModuleWhich()
+                    mapping = []
-                    setattr(sub_module, 'schema', raw_schema)
                     for union_field in field_schema.fields:
-                        setattr(sub_module, union_field.name, union_field.discriminantValue)
+                        mapping.append((union_field.name, union_field.discriminantValue))
+                    sub_module = _StructModuleWhich("StructModuleWhich", mapping)
+                    setattr(sub_module, 'schema', raw_schema)
                 setattr(self, name, sub_module)
         if schema.union_fields and not schema.non_union_fields:
-            sub_module = _StructModuleWhich()
+            mapping = []
             for union_field in schema.node.struct.fields:
                 name = union_field.name
                 name = name[0].upper() + name[1:]
-                setattr(sub_module, name, union_field.discriminantValue)
+                mapping.append((name, union_field.discriminantValue))
+            sub_module = _StructModuleWhich("StructModuleWhich", mapping)
             setattr(self, 'Union', sub_module)
 
     def read(self, file, traversal_limit_in_words=None, nesting_limit=None):

test/test_struct.py

@@ -200,8 +200,13 @@ def test_union_enum(all_types):
 
     msg = all_types.UnionAllTypes.new_message(**{'unionStructField1': {'textField': "foo"}})
     assert msg.which == all_types.UnionAllTypes.Union.UnionStructField1
+    assert msg.which == 'unionStructField1'
+    assert msg.which == 0
+
     msg = all_types.UnionAllTypes.new_message(**{'unionStructField2': {'textField': "foo"}})
     assert msg.which == all_types.UnionAllTypes.Union.UnionStructField2
+    assert msg.which == 'unionStructField2'
+    assert msg.which == 1
 
     assert all_types.GroupedUnionAllTypes.Union.G1 == 0
     assert all_types.GroupedUnionAllTypes.Union.G2 == 1
@@ -212,6 +217,9 @@ def test_union_enum(all_types):
     msg = all_types.GroupedUnionAllTypes.new_message(**{'g2': {'unionStructField2': {'textField': "foo"}}})
     assert msg.which == all_types.GroupedUnionAllTypes.Union.G2
 
+    msg = all_types.UnionAllTypes.new_message()
+    msg.unionStructField2 = msg.init(all_types.UnionAllTypes.Union.UnionStructField2)
+
 
 def isstr(s):
     return isinstance(s, str)