cognite-neat 0.125.1__py3-none-any.whl → 0.126.1__py3-none-any.whl

cognite/neat/_data_model/importers/_table_importer/importer.py

@@ -1,18 +1,25 @@
 from collections.abc import Mapping
+from pathlib import Path
 from typing import ClassVar, cast
 
+import yaml
+from openpyxl import load_workbook
+from openpyxl.worksheet.worksheet import Worksheet
 from pydantic import ValidationError
 
 from cognite.neat._data_model.importers._base import DMSImporter
 from cognite.neat._data_model.models.dms import (
     RequestSchema,
 )
-from cognite.neat._exceptions import DataModelImportError
+from cognite.neat._exceptions import DataModelImportException
 from cognite.neat._issues import ModelSyntaxError
-from cognite.neat._utils.useful_types import CellValueType
+from cognite.neat._utils.text import humanize_collection
+from cognite.neat._utils.useful_types import CellValueType, DataModelTableType
 from cognite.neat._utils.validation import as_json_path, humanize_validation_error
 
-from .data_classes import TableDMS
+from .data_classes import MetadataValue, TableDMS
+from .reader import DMSTableReader
+from .source import SpreadsheetReadContext, TableSource
 
 
 class DMSTableImporter(DMSImporter):
@@ -29,12 +36,45 @@ class DMSTableImporter(DMSImporter):
     REQUIRED_SHEET_MESSAGES: ClassVar[Mapping[str, str]] = {
         f"Missing required column: {sheet!r}": f"Missing required sheet: {sheet!r}" for sheet in REQUIRED_SHEETS
     }
+    MetadataSheet = cast(str, TableDMS.model_fields["metadata"].validation_alias)
 
-    def __init__(self, tables: dict[str, list[dict[str, CellValueType]]]) -> None:
+    def __init__(self, tables: DataModelTableType, source: TableSource | None = None) -> None:
         self._table = tables
+        self._source = source or TableSource("Unknown")
 
     def to_data_model(self) -> RequestSchema:
-        raise NotImplementedError()
+        tables = self._read_tables()
+
+        space, version = self._read_defaults(tables.metadata)
+        reader = DMSTableReader(space, version, self._source)
+        return reader.read_tables(tables)
+
+    @classmethod
+    def from_yaml(cls, yaml_file: Path) -> "DMSTableImporter":
+        """Create a DMSTableImporter from a YAML file."""
+        source = cls._display_name(yaml_file)
+        return cls(yaml.safe_load(yaml_file.read_text()), TableSource(source.as_posix()))
+
+    @classmethod
+    def from_excel(cls, excel_file: Path) -> "DMSTableImporter":
+        """Create a DMSTableImporter from an Excel file."""
+        tables: DataModelTableType = {}
+        source = TableSource(cls._display_name(excel_file).as_posix())
+        workbook = load_workbook(excel_file, read_only=True, data_only=True, rich_text=False)
+        try:
+            for column_id, column_info in TableDMS.model_fields.items():
+                sheet_name = cast(str, column_info.validation_alias)
+                if sheet_name not in workbook.sheetnames:
+                    continue
+                required_columns = TableDMS.get_sheet_columns(column_id, column_info, column_type="required")
+                sheet = workbook[sheet_name]
+                context = SpreadsheetReadContext()
+                table_rows = cls._read_rows(sheet, required_columns, context)
+                tables[sheet_name] = table_rows
+                source.table_read[sheet_name] = context
+            return cls(tables, source)
+        finally:
+            workbook.close()
 
     def _read_tables(self) -> TableDMS:
         try:
@@ -42,7 +82,7 @@ class DMSTableImporter(DMSImporter):
             table = TableDMS.model_validate(self._table)
         except ValidationError as e:
             errors = self._create_error_messages(e)
-            raise DataModelImportError(errors) from None
+            raise DataModelImportException(errors) from None
         return table
 
     def _create_error_messages(self, error: ValidationError) -> list[ModelSyntaxError]:
@@ -74,3 +114,59 @@ class DMSTableImporter(DMSImporter):
             return f"{loc[0]} sheet row {loc[1] + 1} column {loc[2]!r}"
         # This should be unreachable as the TableDMS model only has 2 levels.
         return as_json_path(loc)
+
+    @staticmethod
+    def _read_defaults(metadata: list[MetadataValue]) -> tuple[str, str]:
+        """Reads the space and version from the metadata table."""
+        default_space: str | None = None
+        default_version: str | None = None
+        missing = {"space", "version"}
+        for meta in metadata:
+            if meta.key == "space":
+                default_space = str(meta.value)
+                missing.remove("space")
+            elif meta.key == "version":
+                default_version = str(meta.value)
+                missing.remove("version")
+        if missing:
+            error = ModelSyntaxError(message=f"In Metadata missing required values: {humanize_collection(missing)}")
+            # If space or version is missing, we cannot continue parsing the model as these are used as defaults.
+            raise DataModelImportException([error]) from None
+        return str(default_space), str(default_version)
+
+    @classmethod
+    def _display_name(cls, filepath: Path) -> Path:
+        """Get a display-friendly version of the file path."""
+        cwd = Path.cwd()
+        source = filepath
+        if filepath.is_relative_to(cwd):
+            source = filepath.relative_to(cwd)
+        return source
+
+    @classmethod
+    def _read_rows(
+        cls, sheet: Worksheet, required_columns: list[str], context: SpreadsheetReadContext
+    ) -> list[dict[str, CellValueType]]:
+        table_rows: list[dict[str, CellValueType]] = []
+        # Metadata sheet is just a key-value pair of the first two columns.
+        # For other sheets, we need to find the column header row first.
+        columns: list[str] = [] if sheet.title != cls.MetadataSheet else required_columns
+        for row_no, row in enumerate(sheet.iter_rows(values_only=True)):
+            if columns:
+                # We have found the column header row, read the data rows.
+                if all(cell is None for cell in row):
+                    context.empty_rows.append(row_no)
+                else:
+                    record = dict(zip(columns, row, strict=False))
+                    # MyPy complains as it thinks DataTableFormula | ArrayFormula could be cell values,
+                    # but as we used values_only=True, this is not the case.
+                    table_rows.append(record)  # type: ignore[arg-type]
+            else:
+                # Look for the column header row.
+                row_values = [str(cell) for cell in row]
+                if set(row_values).intersection(required_columns):
+                    columns = row_values
+                    context.header_row = row_no
+                else:
+                    context.skipped_rows.append(row_no)
+        return table_rows