From 0f6df849cd40a5371afa88f62ac3df60a9633d83 Mon Sep 17 00:00:00 2001
From: John Vandenberg <jayvdb@gmail.com>
Date: Tue, 1 Jun 2021 15:40:35 +0800
Subject: [PATCH] Refine exception invoking which on non-union type

Related to https://github.com/capnproto/pycapnp/issues/254
---
 capnp/lib/capnp.pyx | 24 ++++++++++++++++--------
 test/test_struct.py | 22 ++++++++++++++++------
 2 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx
index 42f96c0..f01a63a 100644
--- a/capnp/lib/capnp.pyx
+++ b/capnp/lib/capnp.pyx
@@ -1147,8 +1147,10 @@ cdef class _DynamicStructReader:
     cpdef _which_str(self):
         try:
             return <char *>helpers.fixMaybe(self.thisptr.which()).getProto().getName().cStr()
-        except:
-            raise KjException("Attempted to call which on a non-union type")
+        except RuntimeError as e:
+            if str(e) == "Member was null.":
+                raise KjException("Attempted to call which on a non-union type")
+            raise
 
     cpdef _DynamicEnumField _which(self):
         """Returns the enum corresponding to the union in this struct
@@ -1161,8 +1163,10 @@ cdef class _DynamicStructReader:
         try:
             which = _DynamicEnumField()._init(
                 _StructSchemaField()._init(helpers.fixMaybe(self.thisptr.which()), self).proto)
-        except:
-            raise KjException("Attempted to call which on a non-union type")
+        except RuntimeError as e:
+            if str(e) == "Member was null.":
+                raise KjException("Attempted to call which on a non-union type")
+            raise
 
         return which
 
@@ -1445,8 +1449,10 @@ cdef class _DynamicStructBuilder:
     cpdef _which_str(self):
         try:
             return <char *>helpers.fixMaybe(self.thisptr.which()).getProto().getName().cStr()
-        except:
-            raise KjException("Attempted to call which on a non-union type")
+        except RuntimeError as e:
+            if str(e) == "Member was null.":
+                raise KjException("Attempted to call which on a non-union type")
+            raise
 
     cpdef _DynamicEnumField _which(self):
         """Returns the enum corresponding to the union in this struct
@@ -1459,8 +1465,10 @@ cdef class _DynamicStructBuilder:
         try:
             which = _DynamicEnumField()._init(
                 _StructSchemaField()._init(helpers.fixMaybe(self.thisptr.which()), self).proto)
-        except:
-            raise KjException("Attempted to call which on a non-union type")
+        except RuntimeError as e:
+            if str(e) == "Member was null.":
+                raise KjException("Attempted to call which on a non-union type")
+            raise
 
         return which
 
diff --git a/test/test_struct.py b/test/test_struct.py
index 8dca0fd..153e003 100644
--- a/test/test_struct.py
+++ b/test/test_struct.py
@@ -4,6 +4,8 @@ import os
 import tempfile
 import sys
 
+from capnp.lib.capnp import KjException
+
 this_dir = os.path.dirname(__file__)
 
 
@@ -37,9 +39,13 @@ def test_which_builder(addressbook):
     assert bob.employment.which == addressbook.Person.Employment.unemployed
     assert bob.employment.which == "unemployed"
 
-    with pytest.raises(Exception):
-        addresses.which
-    with pytest.raises(Exception):
+    with pytest.raises(KjException):
+        addresses._which()
+
+    with pytest.raises(KjException):
+        addresses._which_str()
+
+    with pytest.raises(KjException):
         addresses.which
 
 
@@ -71,9 +77,13 @@ def test_which_reader(addressbook):
     bob = people[1]
     assert bob.employment.which == "unemployed"
 
-    with pytest.raises(Exception):
-        addresses.which
-    with pytest.raises(Exception):
+    with pytest.raises(KjException):
+        addresses._which_str()
+
+    with pytest.raises(KjException):
+        addresses._which()
+
+    with pytest.raises(KjException):
         addresses.which