cognite-neat 0.88.3__py3-none-any.whl → 0.90.0__py3-none-any.whl

cognite/neat/rules/importers/_spreadsheet2rules.py

@@ -6,35 +6,30 @@ generating a list of rules based on which nodes that form the graph are made.
 from collections import UserDict, defaultdict
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Literal, cast, overload
+from typing import Literal, cast
 
 import pandas as pd
+from cognite.client.utils._importing import local_import
 from pandas import ExcelFile
 
-from cognite.neat.issues import IssueList, NeatError
+from cognite.neat.issues import IssueList
 from cognite.neat.issues.errors import (
     FileMissingRequiredFieldError,
     FileNotFoundNeatError,
     FileReadError,
     PropertyDefinitionDuplicatedError,
 )
+from cognite.neat.rules._shared import ReadRules, T_InputRules
 from cognite.neat.rules.models import (
-    RULES_PER_ROLE,
-    AssetRules,
-    DMSRules,
-    DomainRules,
-    InformationRules,
+    INPUT_RULES_BY_ROLE,
+    VERIFIED_RULES_BY_ROLE,
     RoleTypes,
     SchemaCompleteness,
 )
-from cognite.neat.rules.models.asset import AssetRulesInput
-from cognite.neat.rules.models.dms import DMSRulesInput
-from cognite.neat.rules.models.information import InformationRulesInput
-from cognite.neat.utils.auxiliary import local_import
 from cognite.neat.utils.spreadsheet import SpreadsheetRead, read_individual_sheet
 from cognite.neat.utils.text import humanize_collection
 
-from ._base import BaseImporter, VerifiedRules, _handle_issues
+from ._base import BaseImporter
 
 SOURCE_SHEET__TARGET_FIELD__HEADERS = [
     (
@@ -53,7 +48,7 @@ SOURCE_SHEET__TARGET_FIELD__HEADERS = [
 ]
 
 MANDATORY_SHEETS_BY_ROLE: dict[RoleTypes, set[str]] = {
-    role_type: {str(sheet_name) for sheet_name in RULES_PER_ROLE[role_type].mandatory_fields(use_alias=True)}
+    role_type: {str(sheet_name) for sheet_name in VERIFIED_RULES_BY_ROLE[role_type].mandatory_fields(use_alias=True)}
     for role_type in RoleTypes.__members__.values()
 }
 
@@ -207,7 +202,7 @@ class SpreadsheetReader:
         return sheets, read_info_by_sheet
 
 
-class ExcelImporter(BaseImporter):
+class ExcelImporter(BaseImporter[T_InputRules]):
     """Import rules from an Excel file.
 
     Args:
@@ -217,30 +212,18 @@ class ExcelImporter(BaseImporter):
     def __init__(self, filepath: Path):
         self.filepath = filepath
 
-    @overload
-    def to_rules(self, errors: Literal["raise"], role: RoleTypes | None = None) -> VerifiedRules: ...
-
-    @overload
-    def to_rules(
-        self,
-        errors: Literal["continue"] = "continue",
-        role: RoleTypes | None = None,
-    ) -> tuple[VerifiedRules | None, IssueList]: ...
-
-    def to_rules(
-        self, errors: Literal["raise", "continue"] = "continue", role: RoleTypes | None = None
-    ) -> tuple[VerifiedRules | None, IssueList] | VerifiedRules:
+    def to_rules(self) -> ReadRules[T_InputRules]:
         issue_list = IssueList(title=f"'{self.filepath.name}'")
         if not self.filepath.exists():
             issue_list.append(FileNotFoundNeatError(self.filepath))
-            return self._return_or_raise(issue_list, errors)
+            return ReadRules(None, issue_list, {})
 
         with pd.ExcelFile(self.filepath) as excel_file:
             user_reader = SpreadsheetReader(issue_list)
 
             user_read = user_reader.read(excel_file, self.filepath)
             if user_read is None or issue_list.has_errors:
-                return self._return_or_raise(issue_list, errors)
+                return ReadRules(None, issue_list, {})
 
             last_read: ReadResult | None = None
             if any(sheet_name.startswith("Last") for sheet_name in user_reader.seen_sheets):
@@ -252,7 +235,7 @@ class ExcelImporter(BaseImporter):
                 reference_read = SpreadsheetReader(issue_list, sheet_prefix="Ref").read(excel_file, self.filepath)
 
         if issue_list.has_errors:
-            return self._return_or_raise(issue_list, errors)
+            return ReadRules(None, issue_list, {})
 
         if reference_read and user_read.role != reference_read.role:
             issue_list.append(
@@ -265,7 +248,7 @@ class ExcelImporter(BaseImporter):
                     "sheet",
                 )
             )
-            return self._return_or_raise(issue_list, errors)
+            return ReadRules(None, issue_list, {})
 
         sheets = user_read.sheets
         original_role = user_read.role
@@ -280,34 +263,12 @@ class ExcelImporter(BaseImporter):
             sheets["reference"] = reference_read.sheets
             read_info_by_sheet.update(reference_read.read_info_by_sheet)
 
-        rules_cls = RULES_PER_ROLE[original_role]
-        with _handle_issues(
-            issue_list,
-            error_cls=NeatError,
-            error_args={"read_info_by_sheet": read_info_by_sheet},
-        ) as future:
-            rules: VerifiedRules
-            if rules_cls is DMSRules:
-                rules = DMSRulesInput.load(sheets).as_rules()
-            elif rules_cls is InformationRules:
-                rules = InformationRulesInput.load(sheets).as_rules()
-            elif rules_cls is AssetRules:
-                rules = AssetRulesInput.load(sheets).as_rules()
-            else:
-                rules = rules_cls.model_validate(sheets)  # type: ignore[attr-defined]
-
-        if future.result == "failure" or issue_list.has_errors:
-            return self._return_or_raise(issue_list, errors)
+        rules_cls = INPUT_RULES_BY_ROLE[original_role]
+        rules = cast(T_InputRules, rules_cls.load(sheets))
+        return ReadRules(rules, issue_list, {"read_info_by_sheet": read_info_by_sheet})
 
-        return self._to_output(
-            rules,
-            issue_list,
-            errors=errors,
-            role=role,
-        )
 
-
-class GoogleSheetImporter(BaseImporter):
+class GoogleSheetImporter(BaseImporter[T_InputRules]):
     """Import rules from a Google Sheet.
 
     .. warning::
@@ -323,38 +284,13 @@ class GoogleSheetImporter(BaseImporter):
         self.sheet_id = sheet_id
         self.skiprows = skiprows
 
-    @overload
-    def to_rules(self, errors: Literal["raise"], role: RoleTypes | None = None) -> VerifiedRules: ...
-
-    @overload
-    def to_rules(
-        self, errors: Literal["continue"] = "continue", role: RoleTypes | None = None
-    ) -> tuple[VerifiedRules | None, IssueList]: ...
+    def to_rules(self) -> ReadRules[T_InputRules]:
+        raise NotImplementedError("Google Sheet Importer is not yet implemented.")
 
-    def to_rules(
-        self, errors: Literal["raise", "continue"] = "continue", role: RoleTypes | None = None
-    ) -> tuple[VerifiedRules | None, IssueList] | VerifiedRules:
+    def _get_sheets(self) -> dict[str, pd.DataFrame]:
         local_import("gspread", "google")
         import gspread  # type: ignore[import]
 
-        role = role or RoleTypes.domain_expert
-        rules_model = cast(DomainRules | InformationRules | AssetRules | DMSRules, RULES_PER_ROLE[role])
-
         client_google = gspread.service_account()
         google_sheet = client_google.open_by_key(self.sheet_id)
-        sheets = {worksheet.title: pd.DataFrame(worksheet.get_all_records()) for worksheet in google_sheet.worksheets()}
-        sheet_names = {str(name).lower() for name in sheets.keys()}
-
-        if missing_sheets := rules_model.mandatory_fields().difference(sheet_names):
-            raise ValueError(f"Missing mandatory sheets: {missing_sheets}")
-
-        if role == RoleTypes.domain_expert:
-            output = rules_model.model_validate(sheets)
-        elif role == RoleTypes.information:
-            output = rules_model.model_validate(sheets)
-        elif role == RoleTypes.dms:
-            output = rules_model.model_validate(sheets)
-        else:
-            raise ValueError(f"Role {role} is not valid.")
-
-        return self._to_output(output, IssueList(), errors=errors, role=role)
+        return {worksheet.title: pd.DataFrame(worksheet.get_all_records()) for worksheet in google_sheet.worksheets()}

cognite/neat/rules/importers/_yaml2rules.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Any, Literal, overload
+from typing import Any, cast
 
 import yaml
 
@@ -11,13 +11,13 @@ from cognite.neat.issues.errors import (
     FileTypeUnexpectedError,
 )
 from cognite.neat.issues.warnings import NeatValueWarning
-from cognite.neat.rules.models import RULES_PER_ROLE, DMSRules, RoleTypes
-from cognite.neat.rules.models.dms import DMSRulesInput
+from cognite.neat.rules._shared import ReadRules, T_InputRules
+from cognite.neat.rules.models import INPUT_RULES_BY_ROLE, RoleTypes
 
-from ._base import BaseImporter, VerifiedRules, _handle_issues
+from ._base import BaseImporter
 
 
-class YAMLImporter(BaseImporter):
+class YAMLImporter(BaseImporter[T_InputRules]):
     """Imports the rules from a YAML file.
 
     Args:
@@ -51,21 +51,9 @@ class YAMLImporter(BaseImporter):
             return cls({}, [FileTypeUnexpectedError(filepath, frozenset([".yaml", ".yml"]))])
         return cls(yaml.safe_load(filepath.read_text()), filepaths=[filepath])
 
-    @overload
-    def to_rules(self, errors: Literal["raise"], role: RoleTypes | None = None) -> VerifiedRules: ...
-
-    @overload
-    def to_rules(
-        self, errors: Literal["continue"] = "continue", role: RoleTypes | None = None
-    ) -> tuple[VerifiedRules | None, IssueList]: ...
-
-    def to_rules(
-        self, errors: Literal["raise", "continue"] = "continue", role: RoleTypes | None = None
-    ) -> tuple[VerifiedRules | None, IssueList] | VerifiedRules:
+    def to_rules(self) -> ReadRules[T_InputRules]:
         if self._read_issues.has_errors or not self.raw_data:
-            if errors == "raise":
-                raise self._read_issues.as_errors()
-            return None, self._read_issues
+            return ReadRules(None, self._read_issues, {})
         issue_list = IssueList(title="YAML Importer", issues=self._read_issues)
 
         if not self._filepaths:
@@ -81,33 +69,18 @@ class YAMLImporter(BaseImporter):
 
         if "metadata" not in self.raw_data:
             self._read_issues.append(FileMissingRequiredFieldError(metadata_file, "section", "metadata"))
-            if errors == "raise":
-                raise self._read_issues.as_errors()
-            return None, self._read_issues
+            return ReadRules(None, self._read_issues, {})
 
         metadata = self.raw_data["metadata"]
 
         if "role" not in metadata:
             self._read_issues.append(FileMissingRequiredFieldError(metadata, "metadata", "role"))
-            if errors == "raise":
-                raise self._read_issues.as_errors()
-            return None, self._read_issues
+            return ReadRules(None, self._read_issues, {})
 
         role_input = RoleTypes(metadata["role"])
         role_enum = RoleTypes(role_input)
-        rules_model = RULES_PER_ROLE[role_enum]
-
-        with _handle_issues(issue_list) as future:
-            rules: VerifiedRules
-            if rules_model is DMSRules:
-                rules = DMSRulesInput.load(self.raw_data).as_rules()
-            else:
-                rules = rules_model.model_validate(self.raw_data)
-
-        if future.result == "failure":
-            if errors == "continue":
-                return None, issue_list
-            else:
-                raise issue_list.as_errors()
-
-        return self._to_output(rules, issue_list, errors, role)
+        rules_cls = INPUT_RULES_BY_ROLE[role_enum]
+
+        rules = cast(T_InputRules, rules_cls.load(self.raw_data))
+
+        return ReadRules(rules, issue_list, {})

cognite/neat/rules/models/__init__.py

@@ -1,12 +1,25 @@
-from cognite.neat.rules.models.asset import AssetRules
-from cognite.neat.rules.models.domain import DomainRules
+from cognite.neat.rules.models.asset._rules import AssetRules
+from cognite.neat.rules.models.asset._rules_input import AssetInputRules
+from cognite.neat.rules.models.domain import DomainInputRules, DomainRules
 from cognite.neat.rules.models.information._rules import InformationRules
+from cognite.neat.rules.models.information._rules_input import InformationInputRules
 
-from ._base import DataModelType, ExtensionCategory, RoleTypes, SchemaCompleteness, SheetEntity, SheetList
+from ._base_rules import DataModelType, ExtensionCategory, RoleTypes, SchemaCompleteness, SheetEntity, SheetList
 from .dms._rules import DMSRules
+from .dms._rules_input import DMSInputRules
 from .dms._schema import DMSSchema
 
-RULES_PER_ROLE: dict[RoleTypes, type[DomainRules] | type[InformationRules] | type[AssetRules] | type[DMSRules]] = {
+INPUT_RULES_BY_ROLE: dict[
+    RoleTypes, type[InformationInputRules] | type[AssetInputRules] | type[DMSInputRules] | type[DomainInputRules]
+] = {
+    RoleTypes.domain_expert: DomainInputRules,
+    RoleTypes.information: InformationInputRules,
+    RoleTypes.asset: AssetInputRules,
+    RoleTypes.dms: DMSInputRules,
+}
+VERIFIED_RULES_BY_ROLE: dict[
+    RoleTypes, type[InformationRules] | type[AssetRules] | type[DMSRules] | type[DomainRules]
+] = {
     RoleTypes.domain_expert: DomainRules,
     RoleTypes.information: InformationRules,
     RoleTypes.asset: AssetRules,
@@ -16,10 +29,13 @@ RULES_PER_ROLE: dict[RoleTypes, type[DomainRules] | type[InformationRules] | typ
 
 __all__ = [
     "DomainRules",
+    "DMSInputRules",
+    "InformationInputRules",
+    "AssetInputRules",
     "InformationRules",
     "AssetRules",
     "DMSRules",
-    "RULES_PER_ROLE",
+    "INPUT_RULES_BY_ROLE",
     "DMSSchema",
     "RoleTypes",
     "SchemaCompleteness",

cognite/neat/rules/models/_base_input.py

@@ -0,0 +1,162 @@
+"""Module for base classes for the input models.
+
+The philosophy of the input models is:
+
+* Provide an easy way to input rules. The type hints are made to be human-friendly, for example, Literal instead of
+  Enum.
+* The .dump() method should fill out defaults and have shortcuts. For example, if the prefix is not provided for
+  a class, then the prefix from the metadata is used. For views, if the class is not provided, it is assumed to
+  be the same as the view.
+
+The base classes are to make it easy to create the input models with default behavior. They are also used for
+testing to ensure that input models correctly map to the verified rules models.
+"""
+
+import sys
+from abc import ABC, abstractmethod
+from dataclasses import Field, dataclass, fields, is_dataclass
+from types import GenericAlias, UnionType
+from typing import Any, Generic, TypeVar, Union, cast, get_args, get_origin, overload
+
+from ._base_rules import BaseRules, RuleModel
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+T_BaseRules = TypeVar("T_BaseRules", bound=BaseRules)
+T_RuleModel = TypeVar("T_RuleModel", bound=RuleModel)
+
+
+@dataclass
+class InputRules(Generic[T_BaseRules], ABC):
+    """Input rules are raw data that is not yet validated."""
+
+    @classmethod
+    @abstractmethod
+    def _get_verified_cls(cls) -> type[T_BaseRules]:
+        raise NotImplementedError("This method should be implemented in the subclass.")
+
+    @classmethod
+    @overload
+    def load(cls, data: dict[str, Any]) -> Self: ...
+
+    @classmethod
+    @overload
+    def load(cls, data: None) -> None: ...
+
+    @classmethod
+    def load(cls, data: dict | None) -> Self | None:
+        if data is None:
+            return None
+        return cls._load(data)
+
+    @classmethod
+    def _type_by_field_name(cls) -> dict[str, type]:
+        output: dict[str, type] = {}
+        for field_ in fields(cls):
+            type_ = field_.type
+            if isinstance(type_, UnionType) or get_origin(type_) is Union:
+                type_ = get_args(type_)[0]
+            if isinstance(type_, str) and type_.startswith(cls.__name__):
+                type_ = cls
+
+            if is_dataclass(type_):
+                candidate = type_
+            elif isinstance(type_, GenericAlias) and type_.__origin__ is list and is_dataclass(type_.__args__[0]):
+                candidate = type_.__args__[0]
+            else:
+                continue
+
+            if hasattr(candidate, "_load"):
+                output[field_.name] = candidate
+        return output
+
+    @classmethod
+    def _load(cls, data: dict[str, Any]) -> Self:
+        args: dict[str, Any] = {}
+        field_type_by_name = cls._type_by_field_name()
+        for field_name, field_ in cls._get_verified_cls().model_fields.items():
+            field_type = field_type_by_name.get(field_name)
+            if field_type is None:
+                continue
+            if field_name in data:
+                value = data[field_name]
+            elif field_.alias in data:
+                value = data[field_.alias]
+            else:
+                continue
+
+            if isinstance(value, dict):
+                args[field_name] = field_type._load(value)  # type: ignore[attr-defined]
+            elif isinstance(value, list) and value and isinstance(value[0], dict):
+                args[field_name] = [field_type._load(item) for item in value]  # type: ignore[attr-defined]
+        return cls(**args)
+
+    def _dataclass_fields(self) -> list[Field]:
+        return list(fields(self))
+
+    def as_rules(self) -> T_BaseRules:
+        cls_ = self._get_verified_cls()
+        return cls_.model_validate(self.dump())
+
+    def dump(self) -> dict[str, Any]:
+        output: dict[str, Any] = {}
+        for field_ in self._dataclass_fields():
+            value = getattr(self, field_.name)
+            if value is None:
+                continue
+            if hasattr(value, "dump"):
+                output[field_.name] = value.dump()
+            elif isinstance(value, list) and value and hasattr(value[0], "dump"):
+                output[field_.name] = [item.dump() for item in value]
+        return output
+
+
+@dataclass
+class InputComponent(ABC, Generic[T_RuleModel]):
+    @classmethod
+    @abstractmethod
+    def _get_verified_cls(cls) -> type[T_RuleModel]:
+        raise NotImplementedError("This method should be implemented in the subclass.")
+
+    @classmethod
+    @overload
+    def load(cls, data: None) -> None: ...
+
+    @classmethod
+    @overload
+    def load(cls, data: dict[str, Any]) -> Self: ...
+
+    @classmethod
+    @overload
+    def load(cls, data: list[dict[str, Any]]) -> list[Self]: ...
+
+    @classmethod
+    def load(cls, data: dict[str, Any] | list[dict[str, Any]] | None) -> Self | list[Self] | None:
+        if data is None:
+            return None
+        if isinstance(data, list) or (isinstance(data, dict) and isinstance(data.get("data"), list)):
+            items = cast(list[dict[str, Any]], data.get("data") if isinstance(data, dict) else data)
+            return [loaded for item in items if (loaded := cls.load(item)) is not None]
+        return cls._load(data)
+
+    @classmethod
+    def _load(cls, data: dict[str, Any]) -> Self:
+        args: dict[str, Any] = {}
+        for field_name, field_ in cls._get_verified_cls().model_fields.items():  # type: ignore[attr-defined]
+            if field_.exclude:
+                continue
+            if field_name in data:
+                args[field_name] = data[field_name]
+            elif field_.alias in data:
+                args[field_name] = data[field_.alias]
+        return cls(**args)
+
+    def dump(self, **kwargs) -> dict[str, Any]:
+        return {
+            field_.alias or name: getattr(self, name)
+            for name, field_ in self._get_verified_cls().model_fields.items()
+            if not field_.exclude
+        }

cognite/neat/rules/models/{_base.py → _base_rules.py}

@@ -10,7 +10,7 @@ import types
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterator
 from functools import wraps
-from typing import Annotated, Any, ClassVar, Generic, Literal, TypeAlias, TypeVar
+from typing import Annotated, Any, ClassVar, Generic, Literal, TypeVar
 
 import pandas as pd
 from pydantic import (
@@ -19,7 +19,6 @@ from pydantic import (
     ConfigDict,
     Field,
     PlainSerializer,
-    constr,
     field_validator,
     model_serializer,
     model_validator,
@@ -36,12 +35,6 @@ else:
 METADATA_VALUE_MAX_LENGTH = 5120
 
 
-def _add_alias(data: dict[str, Any], base_model: type[BaseModel]) -> None:
-    for field_name, field_ in base_model.model_fields.items():
-        if field_name not in data and field_.alias in data:
-            data[field_name] = data[field_.alias]
-
-
 def replace_nan_floats_with_default(values: dict, model_fields: dict[str, FieldInfo]) -> dict:
     output = {}
     for field_name, value in values.items():
@@ -126,10 +119,6 @@ def _get_required_fields(model: type[BaseModel], use_alias: bool = False) -> set
     return required_fields
 
 
-Space: TypeAlias = str
-Description: TypeAlias = constr(min_length=1, max_length=1024)  # type: ignore[valid-type]
-
-
 class SchemaCompleteness(StrEnum):
     complete = "complete"
     partial = "partial"

cognite/neat/rules/models/asset/__init__.py

@@ -1,10 +1,13 @@
 from ._rules import AssetClass, AssetMetadata, AssetProperty, AssetRules
-from ._rules_input import AssetRulesInput
+from ._rules_input import AssetInputClass, AssetInputMetadata, AssetInputProperty, AssetInputRules
 
 __all__ = [
     "AssetRules",
     "AssetMetadata",
     "AssetClass",
     "AssetProperty",
-    "AssetRulesInput",
+    "AssetInputRules",
+    "AssetInputMetadata",
+    "AssetInputClass",
+    "AssetInputProperty",
 ]

cognite/neat/rules/models/asset/_rules.py

@@ -1,13 +1,12 @@
 import sys
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
+from typing import Any, ClassVar, Literal, cast
 
 from pydantic import Field, field_validator, model_validator
 from pydantic.main import IncEx
 from rdflib import Namespace
 
 from cognite.neat.constants import get_default_prefixes
-from cognite.neat.rules.models._base import BaseRules, RoleTypes, SheetList
-from cognite.neat.rules.models.domain import DomainRules
+from cognite.neat.rules.models._base_rules import BaseRules, RoleTypes, SheetList
 from cognite.neat.rules.models.entities import (
     CdfResourceEntityList,
     ClassEntity,
@@ -21,10 +20,6 @@ from cognite.neat.rules.models.information import (
     InformationRules,
 )
 
-if TYPE_CHECKING:
-    from cognite.neat.rules.models.dms._rules import DMSRules
-
-
 if sys.version_info >= (3, 11):
     from typing import Self
 else:
@@ -141,16 +136,3 @@ class AssetRules(BaseRules):
             prefix = self.reference.metadata.prefix
             cleaned[reference] = _AssetRulesSerializer(by_alias, prefix).clean(ref_dump, True)
         return cleaned
-
-    def as_domain_rules(self) -> DomainRules:
-        from ._converter import _AssetRulesConverter
-
-        return _AssetRulesConverter(self.as_information_rules()).as_domain_rules()
-
-    def as_dms_rules(self) -> "DMSRules":
-        from ._converter import _AssetRulesConverter
-
-        return _AssetRulesConverter(self.as_information_rules()).as_dms_rules()
-
-    def as_information_rules(self) -> InformationRules:
-        return InformationRules.model_validate(self.model_dump())