Repository: atviriduomenys/spinta

Issue:
Write `process_complex_content` method
This always goes together with `extension`, i.e. `extension` is inside `complexContent`. 

In addition to this, `complexContent` can have attribute `mixed="true"` (default is `false`) which would mean that this element can also have text. See example in `process_extension` #872 


- [ ] tests
- [ ] documentation
- [ ] entry in CHANGES.rst

Discussion:

Patch:
diff --git a/spinta/manifests/xsd2/helpers.py b/spinta/manifests/xsd2/helpers.py
index ca2b7d5b..552f296a 100644
--- a/spinta/manifests/xsd2/helpers.py
+++ b/spinta/manifests/xsd2/helpers.py
@@ -3,7 +3,8 @@ from __future__ import annotations
 import logging
 import os
 from dataclasses import dataclass
-from typing import Any
+import re
+from typing import Any, List
 from urllib.request import urlopen
 
 from lxml import etree, objectify
@@ -225,6 +226,7 @@ class XSDModel:
     deduplicate_property_name: Deduplicator
     xsd_node_type: str | None = None  # from complexType or from element
     models_by_ref: str | None = None
+    extends_model: XSDModel | None = None
 
     def __init__(self, dataset_resource) -> None:
         self.properties = {}
@@ -340,23 +342,40 @@ class XSDReader:
     def _post_process_refs(self):
         """
         Links properties in all models to their target models based on xsd_ref_to and xsd_type_to.
+        Also links models to their base models based on the 'prepare' attribute (extend statements).
         """
         for model in self.models:
             for prop in model.properties.values():
-                
                 if prop.xsd_ref_to:
                     try:
                         target_model: XSDModel = self.top_level_element_models[prop.xsd_ref_to]
+                        prop.ref_model = target_model
                     except KeyError:
-                        raise KeyError("Reference to a non-existing model")
-                    
+                        raise KeyError(f"Reference to a non-existing model: {prop.xsd_ref_to}")
                 elif prop.xsd_type_to:
                     try:
                         target_model: XSDModel = self.top_level_complex_type_models[prop.xsd_type_to]
+                        prop.ref_model = target_model
+                    except KeyError:
+                        raise KeyError(f"Reference to a non-existing model: {prop.xsd_type_to}")
+
+            if hasattr(model, "prepare") and model.prepare:
+                # Assume the prepare statement is in the format 'extend("BaseType")'
+                prepare_str = model.prepare.strip()
+                model_match: re.Match[str] | None = re.match(r'extend\("(.+)"\)', prepare_str)
+                if model_match:
+                    extends_model_name: str = model_match.group(1)
+                    try:
+                        extends_model: XSDModel | None = self.top_level_complex_type_models[extends_model_name]
                     except KeyError:
-                        raise KeyError("Reference to a non-existing model")
+                        raise KeyError(f"Parent model '{extends_model}' not found for model '{model.name}'")
+                    
+                    extends_props: dict[str, XSDProperty] = extends_model.properties
 
-                prop.ref_model = target_model
+                    if extends_props:
+                        model.extends_model = extends_model
+                    else:
+                        model.prepare = None
 
     def _add_expand_to_top_level_models(self):
         pass
@@ -523,78 +542,91 @@ class XSDReader:
         # todo decide where to deal with placeholder elements, which are not turned into a model
         #  talk to Mantas if a model is considered a placeholder model if it has only one ref or even if it has more refs but nothing else
         # todo handle unique (though it doesn't exist in RC)
-    def process_complex_type(self, node: _Element, state: State) -> list[XSDModel]:
-
+    def process_complex_type(self, node: _Element, state: State) -> List[XSDModel]:
         models = []
         name = node.attrib.get("name")
 
-        properties = []
+        property_groups = [[]]
 
-        # if `mixed="true"` this means, that there can be text inside the element with this complexType, so this means
-        # that we have to add here a property `text` with a source `text()`
-        if node.attrib.get("mixed", False) == "true":
-            property_name = "text"
-            properties.append(XSDProperty(xsd_name=property_name, required=False, source="text()", is_array=False))
-        choice_properties = []
-        sequence_properties = []
+        if node.attrib.get("mixed", "false") == "true":
+            text_prop = XSDProperty(
+                xsd_name="text",
+                required=False,
+                source="text()",
+                is_array=False
+            )
+            for group in property_groups:
+                group.append(text_prop)
 
         for child in node.getchildren():
             # We don't care about comments
             if isinstance(child, etree._Comment):
                 continue
-            if QName(child).localname == "attribute":
-                prop = self.process_attribute(child, state)
-                properties.append(prop)
-            elif QName(child).localname == "sequence":
-                sequence_properties.extend(self.process_sequence(child, state))
-                properties.extend(sequence_properties)
-            elif QName(child).localname == "choice":
-                # technically, there can be only one `choice` node inside `complexType` node.
-                # Two or more are not allowed. If there need to be two or more `choice` nodes,
-                # they need to be inside a `sequence` node
-                choice_properties.extend(self.process_choice(child, state))
-            else:
-                raise RuntimeError(f"This node type cannot be in the complex type: {QName(node).localname}")
 
-        # join direct properties with properties that come from choices
-        properties_groups = []
+            local_name = QName(child).localname
 
-        # if there are properties from `choice`, we
-        if choice_properties:
-            for choice_properties_group in choice_properties:
+            if local_name == "attribute":
+                prop: XSDProperty = self.process_attribute(child, state)
+                for group in property_groups:
+                    group.append(prop)
 
-                choice_properties_group.extend(properties)
+            elif local_name == "sequence":
+                sequence_property_groups: List[List[XSDProperty]] = self.process_sequence(child, state)
+                new_property_groups = []
+                for group in property_groups:
+                    for seq_group in sequence_property_groups:
+                        combined_group = group.copy()
+                        combined_group.extend(seq_group)
+                        new_property_groups.append(combined_group)
+                property_groups = new_property_groups
 
-                properties_groups.append(choice_properties_group)
+            elif local_name == "choice":
+                choice_property_groups: List[List[XSDProperty]] = self.process_choice(child, state)
+                new_property_groups = []
+                for group in property_groups:
+                    for choice_group in choice_property_groups:
+                        combined_group = group.copy()
+                        combined_group.extend(choice_group)
+                        new_property_groups.append(combined_group)
+                property_groups = new_property_groups
 
-        else:
-            properties_groups = [properties]
+            elif local_name == "complexContent":
+                complex_content_property_groups: List[List[XSDModel]] = self.process_complex_content(child, state)
+                new_property_groups = []
+                for group in property_groups:
+                    for cc_group in complex_content_property_groups:
+                        combined_group = group.copy()
+                        combined_group.extend(cc_group)
+                        new_property_groups.append(combined_group)
+                property_groups = new_property_groups
 
-        new_properties_groups = []
+                if hasattr(state, 'prepare_statement'):
+                    prepare_statement: str = state.prepare_statement
+                    del state.prepare_statement
+                else:
+                    prepare_statement = None
 
-        if sequence_properties:
-            for sequence_properties_group in sequence_properties:
-                for group in properties_groups:
-                    group.extend(sequence_properties_group)
-                    new_properties_groups.append(group)
-        else:
-            new_properties_groups = properties_groups
+            else:
+                raise RuntimeError(f"This node type cannot be in the complex type: {local_name}")
 
-        for group in new_properties_groups:
-            # we don't set model source here, we do it in the process_element or nowhere if it's top level
+        for group in property_groups:
             model = XSDModel(dataset_resource=self.dataset_resource)
             property_deduplicate = Deduplicator()
             for prop in group:
                 prop.name = property_deduplicate(prop.xsd_name)
                 model.properties[prop.name] = prop
+
             if name:
                 model.xsd_name = name
                 model.name = self.deduplicate_model_name(to_model_name(name))
                 self.top_level_complex_type_models[model.xsd_name] = model
+
+            if 'prepare_statement' in locals() and prepare_statement:
+                model.prepare = prepare_statement
+
             models.append(model)
 
         self.models.extend(models)
-
         return models
 
     def _map_type(self, xsd_type: str) -> XSDType:
@@ -740,8 +772,51 @@ class XSDReader:
     def process_simple_content(self, node: _Element, state: State) -> None:
         pass
 
-    def process_complex_content(self, node: _Element, state: State) -> None:
-        pass
+    def process_complex_content(self, node: _Element, state: State) -> list[list[XSDModel]]:
+        property_groups: list[list] = [[]]
+
+        for child in node:
+            if isinstance(child, etree._Comment):
+                continue
+
+            local_name = QName(child).localname
+
+            if local_name == "extension":
+                extension_property_groups: List[List[XSDProperty]] = self.process_complex_type_extension(child, state)
+                new_property_groups = []
+                for group in property_groups:
+                    for ext_group in extension_property_groups:
+                        combined_group = group.copy()
+                        combined_group.extend(ext_group)
+                        new_property_groups.append(combined_group)
+                property_groups = new_property_groups
+
+            elif local_name == "restriction":
+                restriction_prop_type: XSDType = self.process_restriction(child, state)
+                prop_name = node.attrib.get("name", "value")
+                restriction_prop = XSDProperty(
+                    xsd_name=prop_name,
+                    property_type=restriction_prop_type,
+                )
+                for group in property_groups:
+                    group.append(restriction_prop)
+
+            else:
+                raise RuntimeError(f"Unexpected element '{local_name}' in complexContent")
+
+        if node.attrib.get("mixed", "false") == "true":
+            property_name = "text"
+            text_prop = XSDProperty(
+                xsd_name=property_name,
+                property_type=XSDType(name="string"),
+                required=False,
+                source="text()",
+                is_array=False
+            )
+            for group in property_groups:
+                group.append(text_prop)
+
+        return property_groups
 
     def process_annotation(self, node: _Element, state: State) -> str:
         description = ""
@@ -757,9 +832,84 @@ class XSDReader:
 
     def process_documentation(self, node: _Element, state: State) -> str:
         return node.text
+    
+    def process_simple_type_extension(self, node: _Element, state: State) -> tuple[str, list] | None:
+        # this is an initial implementation, this method needs to be finished together with process_simple_type
+        base = node.attrib.get("base")
+        if not base:
+            raise RuntimeError("Extension must have a 'base' attribute.")
+
+        base_type_name = base.split(":")[-1]
+
+        if base_type_name in DATATYPES_MAPPING:
+            type_name = DATATYPES_MAPPING[base_type_name]
+            attributes = []
+
+            for child in node:
+                if isinstance(child, etree._Comment):
+                    continue
+                local_name = QName(child).localname
+                if local_name == "attribute":
+                    prop = self.process_attribute(child, state)
+                    attributes.append(prop)
+                elif local_name == "annotation":
+                    prop = self.process_annotation(child, state)
+                    attributes.append(prop)
+                else:
+                    raise RuntimeError(f"Unexpected element '{local_name}' in simpleType extension")
+
+            return type_name, attributes
+
+    def process_complex_type_extension(self, node: _Element, state: State) -> List[List[XSDProperty]]:
+        base = node.attrib.get("base")
+        if not base:
+            raise RuntimeError("Extension must have a 'base' attribute.")
 
-    def process_extension(self, node: _Element, state: State) -> None:
-        pass
+        base_type_name = base.split(":")[-1]
+
+        base_model: XSDModel | None = self.top_level_complex_type_models.get(base_type_name)
+        if not base_model:
+            raise RuntimeError(f"Base type '{base}' not found.")
+
+        property_groups: list[list] = [[]]
+
+        for child in node.getchildren():
+            if isinstance(child, etree._Comment):
+                continue
+
+            local_name = QName(child).localname
+
+            if local_name == "element":
+                prop: List[XSDProperty] = self.process_element(child, state)
+                for group in property_groups:
+                    group.append(prop)
+
+            elif local_name == "sequence":
+                sequence_property_groups: List[List[XSDProperty]] = self.process_sequence(child, state)
+                new_property_groups = []
+                for group in property_groups:
+                    for seq_group in sequence_property_groups:
+                        combined_group = group.copy()
+                        combined_group.extend(seq_group)
+                        new_property_groups.append(combined_group)
+                property_groups = new_property_groups
+    
+            elif local_name == "choice":
+                choice_property_groups: List[List[XSDProperty]] = self.process_choice(child, state)
+                new_property_groups = []
+                for group in property_groups:
+                    for choice_group in choice_property_groups:
+                        combined_group = group.copy()
+                        combined_group.extend(choice_group)
+                        new_property_groups.append(combined_group)
+                property_groups = new_property_groups
+
+            else:
+                raise RuntimeError(f"Unexpected element '{local_name}' in extension")
+
+        state.prepare_statement = f'extend("{base_type_name}")'
+
+        return property_groups
 
     def process_restriction(self, node: _Element, state: State) -> XSDType:
 


Test patch:
diff --git a/tests/manifests/xsd2/test_helpers.py b/tests/manifests/xsd2/test_helpers.py
index dab00c5c..b1aa60b8 100644
--- a/tests/manifests/xsd2/test_helpers.py
+++ b/tests/manifests/xsd2/test_helpers.py
@@ -339,6 +339,39 @@ def test_process_element_complex_type_separate():
     # Assert that 'required' is True
     assert result.required is True
 
+def test_process_complex_type_with_extension(xsd_reader, create_xsd_model):
+    complex_type_xml = """
+    <complexType name="DerivedType">
+        <complexContent>
+            <extension base="BaseType">
+                <sequence>
+                    <element name="extendedProp" type="xs:string"/>
+                </sequence>
+                <element name="attr1" type="xs:int"/>
+            </extension>
+        </complexContent>
+    </complexType>
+    """
+    complex_type_node = etree.fromstring(complex_type_xml)
+
+    xsd_reader.process_complex_content = MagicMock(return_value=[[
+        XSDProperty(xsd_name="extendedProp", property_type=XSDType(name="string")),
+        XSDProperty(xsd_name="attr1", property_type=XSDType(name="int")),
+    ]])
+    state = State()
+    state.prepare_statement = 'extend("BaseType")'
+    base_model = create_xsd_model("BaseType")
+    xsd_reader.top_level_complex_type_models["BaseType"] = base_model
+
+    models = xsd_reader.process_complex_type(complex_type_node, state)
+
+    assert len(models) == 1, "Should return one model"
+    model = models[0]
+    prop_names = list(model.properties.keys())
+
+    assert "extendedProp" in prop_names
+    assert "attr1" in prop_names
+    assert model.prepare == 'extend("BaseType")'
 
 def test_process_element_ref():
     xsd_schema = """
@@ -796,4 +829,259 @@ def test_post_process_refs_links_existing_references(setup_models):
     
     # Property2 in ModelB should link to ModelA via ref_model
     assert prop2.ref_model is not None
-    assert prop2.ref_model == model_a
\ No newline at end of file
+    assert prop2.ref_model == model_a
+
+def test_post_process_refs_valid_prepare_with_properties(xsd_reader, create_xsd_model):
+
+    extends_model = create_xsd_model("BaseType")
+    extends_model.properties = {
+        "baseProp": XSDProperty(xsd_name="baseProp", property_type=XSDType(name="string"))
+    }
+    
+    derived_model = create_xsd_model("DerivedType")
+    derived_model.prepare = 'extend("BaseType")'
+    derived_model.properties = {
+        "derivedProp": XSDProperty(xsd_name="derivedProp", property_type=XSDType(name="int"))
+    }
+    
+    xsd_reader.models = [derived_model]
+    xsd_reader.top_level_complex_type_models = {"BaseType": extends_model}
+    
+    xsd_reader._post_process_refs()
+    
+    assert derived_model.prepare == 'extend("BaseType")'
+    assert derived_model.extends_model is extends_model
+
+def test_post_process_refs_valid_prepare_with_empty_properties(xsd_reader, create_xsd_model):
+    extends_model = create_xsd_model("BaseType")
+    
+    derived_model = create_xsd_model("DerivedType")
+    derived_model.prepare = 'extend("BaseType")'
+    derived_model.properties = {
+        "derivedProp": XSDProperty(xsd_name="derivedProp", property_type=XSDType(name="int"))
+    }
+
+    xsd_reader.models = [derived_model]
+    xsd_reader.top_level_complex_type_models = {"BaseType": extends_model}
+    
+    xsd_reader._post_process_refs()
+    
+    assert derived_model.prepare is None
+    assert derived_model.extends_model is None
+
+# def test_process_extension_simple_type(xsd_reader):
+#     extension_xml = """
+#     <extension base="xs:string">
+#         <!-- No child elements -->
+#     </extension>
+#     """
+#     extension_node = etree.fromstring(extension_xml)
+
+#     state = State()
+
+#     type_name = xsd_reader.process_extension(extension_node, state)
+
+#     assert type_name == "string", "Should return the type name 'string'"
+
+def test_process_extension_complex_type_no_children(xsd_reader, create_xsd_model):
+    extension_xml = """
+    <extension base="BaseType">
+        <!-- No child elements -->
+    </extension>
+    """
+    extension_node = etree.fromstring(extension_xml)
+
+    state = State()
+    base_model = create_xsd_model("BaseType")
+    xsd_reader.top_level_complex_type_models["BaseType"] = base_model
+
+    property_groups = xsd_reader.process_complex_type_extension(extension_node, state)
+
+    assert isinstance(property_groups, list)
+    assert len(property_groups) == 1
+    assert property_groups[0] == []
+    assert state.prepare_statement == 'extend("BaseType")'
+
+def test_process_extension_complex_type_with_sequence(xsd_reader, create_xsd_model):
+    extension_xml = """
+    <extension base="BaseType">
+        <sequence>
+            <element name="newProp1" type="xs:string"/>
+            <element name="newProp2" type="xs:int"/>
+        </sequence>
+    </extension>
+    """
+    extension_node = etree.fromstring(extension_xml)
+
+    xsd_reader.process_sequence = MagicMock(return_value=[
+        [XSDProperty(xsd_name="newProp1", property_type=XSDType(name="string")),
+         XSDProperty(xsd_name="newProp2", property_type=XSDType(name="int"))]
+    ])
+
+    state = State()
+    base_model = create_xsd_model("BaseType")
+    xsd_reader.top_level_complex_type_models["BaseType"] = base_model
+
+    property_groups = xsd_reader.process_complex_type_extension(extension_node, state)
+
+    assert len(property_groups) == 1
+    prop_names = [prop.xsd_name for prop in property_groups[0]]
+    assert "newProp1" in prop_names
+    assert "newProp2" in prop_names
+    assert state.prepare_statement == 'extend("BaseType")'
+
+def test_process_extension_complex_type_with_choice(xsd_reader, create_xsd_model):
+    extension_xml = """
+    <extension base="BaseType">
+        <choice>
+            <element name="choiceProp1" type="xs:string"/>
+            <element name="choiceProp2" type="xs:int"/>
+        </choice>
+    </extension>
+    """
+    extension_node = etree.fromstring(extension_xml)
+
+    xsd_reader.process_choice = MagicMock(return_value=[
+        [XSDProperty(xsd_name="choiceProp1", property_type=XSDType(name="string"))],
+        [XSDProperty(xsd_name="choiceProp2", property_type=XSDType(name="int"))],
+    ])
+
+    state = State()
+    base_model = create_xsd_model("BaseType")
+    xsd_reader.top_level_complex_type_models["BaseType"] = base_model
+
+    property_groups = xsd_reader.process_complex_type_extension(extension_node, state)
+
+    assert len(property_groups) == 2
+    prop_names_group1 = [prop.xsd_name for prop in property_groups[0]]
+    prop_names_group2 = [prop.xsd_name for prop in property_groups[1]]
+    assert "choiceProp1" in prop_names_group1
+    assert "choiceProp2" not in prop_names_group1
+    assert "choiceProp2" in prop_names_group2
+    assert "choiceProp1" not in prop_names_group2
+    assert state.prepare_statement == 'extend("BaseType")'
+
+def test_process_extension_complex_type_with_elements(xsd_reader, create_xsd_model):
+    extension_xml = """
+    <extension base="BaseType">
+        <element name="attr1" type="xs:string"/>
+        <element name="attr2" type="xs:int"/>
+    </extension>
+    """
+    extension_node = etree.fromstring(extension_xml)
+
+    def mock_process_element(node, state):
+        name = node.attrib.get("name")
+        type_name = node.attrib.get("type").split(":")[-1]
+        return XSDProperty(xsd_name=name, property_type=XSDType(name=type_name))
+    xsd_reader.process_element = mock_process_element
+
+    state = State()
+    base_model = create_xsd_model("BaseType")
+    xsd_reader.top_level_complex_type_models["BaseType"] = base_model
+
+    property_groups = xsd_reader.process_complex_type_extension(extension_node, state)
+
+    assert len(property_groups) == 1
+    prop_names = [prop.xsd_name for prop in property_groups[0]]
+    assert "attr1" in prop_names
+    assert "attr2" in prop_names
+    assert state.prepare_statement == 'extend("BaseType")'
+
+def test_process_complex_content_with_extension(xsd_reader, create_xsd_model):
+    complex_content_xml = """
+    <complexContent>
+        <extension base="BaseType">
+            <sequence>
+                <element name="seqProp1" type="xs:string"/>
+            </sequence>
+            <choice>
+                <element name="choiceProp1" type="xs:int"/>
+                <element name="choiceProp2" type="xs:float"/>
+            </choice>
+            <element name="attr1" type="xs:boolean"/>
+        </extension>
+    </complexContent>
+    """
+    complex_content_node = etree.fromstring(complex_content_xml)
+
+    state = State()
+    base_model = create_xsd_model("BaseType")
+    xsd_reader.top_level_complex_type_models["BaseType"] = base_model
+
+    def mock_process_complex_type_extension(node, state) -> list[list[XSDProperty]]:
+        state.prepare_statement = 'extend("BaseType")'
+        return [
+            [
+                XSDProperty(xsd_name="seqProp1", property_type=XSDType(name="string")),
+                XSDProperty(xsd_name="choiceProp1", property_type=XSDType(name="int")),
+                XSDProperty(xsd_name="attr1", property_type=XSDType(name="boolean")),
+            ],
+            [
+                XSDProperty(xsd_name="seqProp1", property_type=XSDType(name="string")),
+                XSDProperty(xsd_name="choiceProp2", property_type=XSDType(name="float")),
+                XSDProperty(xsd_name="attr1", property_type=XSDType(name="boolean")),
+            ],
+        ]
+    xsd_reader.process_complex_type_extension = mock_process_complex_type_extension
+
+    property_groups = xsd_reader.process_complex_content(complex_content_node, state)
+
+    assert len(property_groups) == 2
+    prop_names_group1 = [prop.xsd_name for prop in property_groups[0]]
+    prop_names_group2 = [prop.xsd_name for prop in property_groups[1]]
+    assert "seqProp1" in prop_names_group1 and "seqProp1" in prop_names_group2
+    assert "choiceProp1" in prop_names_group1 and "choiceProp1" not in prop_names_group2
+    assert "choiceProp2" in prop_names_group2 and "choiceProp2" not in prop_names_group1
+    assert "attr1" in prop_names_group1 and "attr1" in prop_names_group2
+    assert state.prepare_statement == 'extend("BaseType")'
+
+def test_process_complex_content_with_restriction(xsd_reader):
+    complex_content_xml = """
+    <complexContent>
+        <restriction base="xs:string">
+            <enumeration value="Option1"/>
+            <enumeration value="Option2"/>
+        </restriction>
+    </complexContent>
+    """
+    complex_content_node = etree.fromstring(complex_content_xml)
+
+    def mock_process_restriction(node, state):
+        prop_type = XSDType(name="string")
+        prop_type.enums = {"Option1": None, "Option2": None}
+        return prop_type
+    xsd_reader.process_restriction = mock_process_restriction
+
+    state = State()
+
+    property_groups = xsd_reader.process_complex_content(complex_content_node, state)
+
+    assert len(property_groups) == 1
+    prop_group = property_groups[0]
+    assert len(prop_group) == 1
+    prop = prop_group[0]
+    assert prop.xsd_name == "value"
+    assert prop.type.name == "string"
+    assert prop.type.enums == {"Option1": None, "Option2": None}
+
+def test_process_complex_content_with_mixed_content(xsd_reader):
+    complex_content_xml = """
+    <complexContent mixed="true">
+        <extension base="BaseType">
+            <!-- No child elements -->
+        </extension>
+    </complexContent>
+    """
+    complex_content_node = etree.fromstring(complex_content_xml)
+
+    xsd_reader.process_complex_type_extension = MagicMock(return_value=[[]])
+
+    state = State()
+
+    property_groups = xsd_reader.process_complex_content(complex_content_node, state)
+
+    assert len(property_groups) == 1
+    prop_group = property_groups[0]
+    prop_names = [prop.xsd_name for prop in prop_group]
+    assert "text" in prop_names