schema - extending a discriminated union yields new classes

Author: commonism

src/aiopenapi3/openapi.py

@@ -627,6 +627,16 @@ def _init_schema_types(self, only_required: bool) -> None:
             for j in b._model_types:
                 types[j.__name__] = j
 
+        # as previous .get_type() may have created new models, we need to reindex
+        for name, schema in list(types.items()):
+            if not is_basemodel(schema):
+                continue
+            thes = byname.get(name, None)
+            if thes is not None:
+                for v in byid[id(thes)]._model_types:
+                    if v.__name__ not in types:
+                        types[v.__name__] = v
+
         # print(f"{len(types)}")
         for name, schema in types.items():
             if not is_basemodel(schema):
@@ -637,6 +647,7 @@ def _init_schema_types(self, only_required: bool) -> None:
                 thes = byname.get(name, None)
                 if thes is not None:
                     for v in byid[id(thes)]._model_types:
+                        assert v.__name__ in types, v.__name__
                         v.model_rebuild(_types_namespace={"__types": types})
             except Exception as e:
                 raise e

tests/conftest.py

@@ -566,3 +566,8 @@ def with_schema_anyOf():
 @pytest.fixture
 def with_schema_title_name_collision():
     yield _get_parsed_yaml("schema-title-name-collision.yaml")
+
+
+@pytest.fixture
+def with_schema_discriminated_union_extends():
+    yield _get_parsed_yaml("schema-discriminated-union-extends.yaml")

tests/fixtures/schema-discriminated-union-extends.yaml

@@ -0,0 +1,42 @@
+openapi: 3.1.0
+
+info:
+  title: title collision
+  version: 2.2.7
+
+servers:
+  - url: /
+
+paths: {}
+
+
+components:
+  schemas:
+    A:
+      type: object
+      additionalProperties: false
+      properties:
+        a:
+          type: integer
+    B:
+      type: object
+      additionalProperties: false
+      properties:
+        b:
+          type: integer
+
+    AB:
+      type: object
+      additionalProperties: false
+      properties:
+        type:
+          type: string
+          const: ""
+      oneOf:
+        - $ref: "#/components/schemas/A"
+        - $ref: "#/components/schemas/B"
+      discriminator:
+        propertyName: type
+        mapping:
+          A: "#/components/schemas/A"
+          B: "#/components/schemas/B"