Fix a bug where nD entity classes failed to import (#404)

spine-tools · May 14, 2024 · bf5ead7 · bf5ead7
2 parents fc4b515 + e2aea79
commit bf5ead7
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 3 deletions.
diff --git a/spinedb_api/import_mapping/import_mapping.py b/spinedb_api/import_mapping/import_mapping.py
@@ -454,10 +454,13 @@ class DimensionMapping(ImportMapping):
     MAP_TYPE = "Dimension"
 
     def _import_row(self, source_data, state, mapped_data):
-        _ = state[ImportKey.ENTITY_CLASS_NAME]
-        dimension_name = str(source_data)
-        state[ImportKey.DIMENSION_NAMES].append(dimension_name)
+        if ImportKey.ENTITY_CLASS_NAME not in state:
+            raise KeyError(ImportKey.ENTITY_CLASS_NAME)
         dimension_names = state[ImportKey.DIMENSION_NAMES]
+        if len(dimension_names) == state[ImportKey.DIMENSION_COUNT]:
+            return
+        dimension_name = str(source_data)
+        dimension_names.append(dimension_name)
         if len(dimension_names) == state[ImportKey.DIMENSION_COUNT]:
             raise KeyFix(ImportKey.DIMENSION_NAMES)
 

diff --git a/tests/import_mapping/test_generator.py b/tests/import_mapping/test_generator.py
@@ -458,6 +458,61 @@ def test_import_scenario(self):
             {"scenarios": {("scen1",)}},
         )
 
+    def test_importing_multidimensional_class_when_there_is_an_extra_column(self):
+        header = ["3D entity class", "unit", "node", "node", "parameter", "alternative", "value", None]
+        data_source = iter(
+            [
+                ["unit__node__node", "Dyson sphere", "Gamma Ceti", "Ring world", "flow", "Base", 23.3, None],
+                7 * [None] + ["aa"],
+            ]
+        )
+        mappings = [
+            [
+                {"map_type": "EntityClass", "position": 0},
+                {"map_type": "Dimension", "position": "header", "value": 1},
+                {"map_type": "Dimension", "position": "header", "value": 2},
+                {"map_type": "Dimension", "position": "header", "value": 3},
+                {"map_type": "Entity", "position": "hidden"},
+                {"map_type": "Element", "position": 1, "import_entities": True},
+                {"map_type": "Element", "position": 2, "import_entities": True},
+                {"map_type": "Element", "position": 3, "import_entities": True},
+                {"map_type": "EntityMetadata", "position": "hidden"},
+                {"map_type": "ParameterDefinition", "position": 4},
+                {"map_type": "Alternative", "position": 5},
+                {"map_type": "ParameterValueMetadata", "position": "hidden"},
+                {"map_type": "ParameterValue", "position": 6},
+            ]
+        ]
+        convert_function_specs = {
+            0: "string",
+            1: "string",
+            2: "string",
+            3: "string",
+            4: "string",
+            5: "string",
+            6: "float",
+        }
+        convert_functions = {column: value_to_convert_spec(spec) for column, spec in convert_function_specs.items()}
+        mapped_data, errors = get_mapped_data(data_source, mappings, header, column_convert_fns=convert_functions)
+        self.assertEqual(errors, [])
+        self.assertEqual(
+            mapped_data,
+            {
+                "alternatives": {"Base"},
+                "entities": [
+                    ("unit", "Dyson sphere"),
+                    ("node", "Gamma Ceti"),
+                    ("node", "Ring world"),
+                    ("unit__node__node", ("Dyson sphere", "Gamma Ceti", "Ring world")),
+                ],
+                "entity_classes": [("unit",), ("node",), ("unit__node__node", ("unit", "node", "node"))],
+                "parameter_definitions": [("unit__node__node", "flow")],
+                "parameter_values": [
+                    ["unit__node__node", ("Dyson sphere", "Gamma Ceti", "Ring world"), "flow", 23.3, "Base"]
+                ],
+            },
+        )
+
 
 if __name__ == "__main__":
     unittest.main()