ClassInnersHandler: evaluate attr choices as well

tefra · Jun 19, 2021 · 1a2b62a · 1a2b62a
1 parent b71f0a4
commit 1a2b62a
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 0 deletions.
diff --git a/tests/codegen/handlers/test_class_inners.py b/tests/codegen/handlers/test_class_inners.py
@@ -1,6 +1,9 @@
+from typing import Generator
+
 from xsdata.codegen.handlers import ClassInnersHandler
 from xsdata.models.enums import DataType
 from xsdata.utils.testing import AttrFactory
+from xsdata.utils.testing import AttrTypeFactory
 from xsdata.utils.testing import ClassFactory
 from xsdata.utils.testing import ExtensionFactory
 from xsdata.utils.testing import FactoryTestCase
@@ -71,3 +74,37 @@ def test_rename_inner(self):
 
         self.assertEqual("{xsdata}foo_Inner", outer.attrs[0].types[0].qname)
         self.assertEqual("{xsdata}foo_Inner", outer.inner[0].qname)
+
+    def test_find_attr_types_with_attr_choices(self):
+        choices = [
+            AttrFactory.create(
+                types=[
+                    AttrTypeFactory.create("bar", forward=True),
+                    AttrTypeFactory.create("foo", forward=True),
+                ]
+            ),
+            AttrFactory.reference("foo"),
+            AttrFactory.reference("foo", forward=True),
+            AttrFactory.reference("bar", forward=True),
+        ]
+        choice = AttrFactory.create(
+            name="attr_B_Or_attr_C",
+            tag="Choice",
+            index=0,
+            types=[AttrTypeFactory.native(DataType.ANY_TYPE)],
+            choices=choices,
+        )
+        target = ClassFactory.create()
+        target.attrs.append(choice)
+
+        result = self.processor.find_attr_types(target, "foo")
+        self.assertIsInstance(result, Generator)
+
+        self.assertEqual(choices[0].types[1], next(result))
+        self.assertEqual(choices[2].types[0], next(result))
+        self.assertIsNone(next(result, None))
+
+        result = self.processor.find_attr_types(target, "bar")
+        self.assertEqual(choices[0].types[0], next(result))
+        self.assertEqual(choices[3].types[0], next(result))
+        self.assertIsNone(next(result, None))
diff --git a/xsdata/codegen/handlers/class_inners.py b/xsdata/codegen/handlers/class_inners.py
@@ -64,3 +64,8 @@ def find_attr_types(cls, target: Class, qname: str) -> Iterator[AttrType]:
             for attr_type in attr.types:
                 if attr_type.forward and attr_type.qname == qname:
                     yield attr_type
+
+            for choice in attr.choices:
+                for choice_type in choice.types:
+                    if choice_type.forward and choice_type.qname == qname:
+                        yield choice_type