From 0c904443bd227d2de3fb47b438c75712c48b49de Mon Sep 17 00:00:00 2001
From: John Vandenberg <jayvdb@gmail.com>
Date: Mon, 31 May 2021 16:14:21 +0800
Subject: [PATCH] Add Union on top level union messages

Closes https://github.com/capnproto/pycapnp/issues/247
---
 capnp/lib/capnp.pyx  |  7 +++++++
 test/all_types.capnp | 18 ++++++++++++++++++
 test/test_struct.py  | 21 +++++++++++++++++++++
 3 files changed, 46 insertions(+)

diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx
index 7f26575..c9c211c 100644
--- a/capnp/lib/capnp.pyx
+++ b/capnp/lib/capnp.pyx
@@ -3175,6 +3175,13 @@ class _StructModule(object):
                     for union_field in field_schema.fields:
                         setattr(sub_module, union_field.name, union_field.discriminantValue)
                 setattr(self, name, sub_module)
+        if schema.union_fields and not schema.non_union_fields:
+            sub_module = _StructModuleWhich()
+            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)
+            setattr(self, 'Union', sub_module)
 
     def read(self, file, traversal_limit_in_words=None, nesting_limit=None):
         """Returns a Reader for the unpacked object read from file.
diff --git a/test/all_types.capnp b/test/all_types.capnp
index 7dd1550..2984bca 100644
--- a/test/all_types.capnp
+++ b/test/all_types.capnp
@@ -50,3 +50,21 @@ struct TestAllTypes {
   enumList      @32 : List(TestEnum);
   interfaceList @33 : List(Void);  # TODO
 }
+
+struct UnionAllTypes {
+  union {
+    unionStructField1 @0 : TestAllTypes;
+    unionStructField2 @1 : TestAllTypes;
+  }
+}
+
+struct GroupedUnionAllTypes {
+  union {
+    g1 :group {
+      unionStructField1 @0 : TestAllTypes;
+    }
+    g2 :group {
+      unionStructField2 @1 : TestAllTypes;
+    }
+  }
+}
diff --git a/test/test_struct.py b/test/test_struct.py
index 3548011..17655d7 100644
--- a/test/test_struct.py
+++ b/test/test_struct.py
@@ -189,9 +189,30 @@ def test_set_dict(all_types):
 def test_set_dict_union(addressbook):
     person = addressbook.Person.new_message(**{'employment': {'employer': {'name': 'foo'}}})
 
+    assert person.employment.which == addressbook.Person.Employment.employer
+
     assert person.employment.employer.name == 'foo'
 
 
+def test_union_enum(all_types):
+    assert all_types.UnionAllTypes.Union.UnionStructField1 == 0
+    assert all_types.UnionAllTypes.Union.UnionStructField2 == 1
+
+    msg = all_types.UnionAllTypes.new_message(**{'unionStructField1': {'textField': "foo"}})
+    assert msg.which == all_types.UnionAllTypes.Union.UnionStructField1
+    msg = all_types.UnionAllTypes.new_message(**{'unionStructField2': {'textField': "foo"}})
+    assert msg.which == all_types.UnionAllTypes.Union.UnionStructField2
+
+    assert all_types.GroupedUnionAllTypes.Union.G1 == 0
+    assert all_types.GroupedUnionAllTypes.Union.G2 == 1
+
+    msg = all_types.GroupedUnionAllTypes.new_message(**{'g1': {'unionStructField1': {'textField': "foo"}}})
+    assert msg.which == all_types.GroupedUnionAllTypes.Union.G1
+
+    msg = all_types.GroupedUnionAllTypes.new_message(**{'g2': {'unionStructField2': {'textField': "foo"}}})
+    assert msg.which == all_types.GroupedUnionAllTypes.Union.G2
+
+
 def isstr(s):
     return isinstance(s, str)