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

cognite/neat/rules/importers/_rdf/_imf2rules/_imf2classes.py

@@ -37,16 +37,20 @@ def parse_imf_to_classes(graph: Graph, language: str = "en") -> list[dict]:
     query = """
     SELECT ?class ?name ?description ?parentClass ?reference ?match ?comment
     WHERE {
-        #Finding IMF - elements
+        # Finding IMF - elements
         VALUES ?type { imf:BlockType imf:TerminalType imf:AttributeType }
         ?imfClass a ?type .
         OPTIONAL {?imfClass rdfs:subClassOf ?parent }.
         OPTIONAL {?imfClass rdfs:label | skos:prefLabel ?name }.
-        OPTIONAL {?imfClass rdfs:comment | skos:description ?description} .
+
+        # Note: Bug in PCA has lead to the use non-existing term skos:description. This will be replaced
+        # with the correct skos:definition in the near future, so both terms are included here.
+        OPTIONAL {?imfClass rdfs:comment | skos:definition | skos:description ?description} .
 
         # Finding the last segment of the class IRI
         BIND(STR(?imfClass) AS ?classString)
-        BIND(REPLACE(?classString, "^.*[/#]([^/#]*)$", "$1") AS ?classSegment)
+        BIND(REPLACE(?classString, "^.*[/#]([^/#]*)$", "$1") AS ?tempSegment)
+        BIND(REPLACE(?tempSegment, "-", "_") AS ?classSegment)
         BIND(IF(CONTAINS(?classString, "imf/"), CONCAT("IMF_", ?classSegment) , ?classSegment) AS ?class)
 
         # Add imf:Attribute as parent class

cognite/neat/rules/importers/_rdf/_imf2rules/_imf2metadata.py

@@ -19,12 +19,12 @@ def parse_imf_metadata() -> dict:
     raw_metadata = {
         "role": RoleTypes.information,
         "schema": SchemaCompleteness.partial,
-        "prefix": "pca-imf",
-        "namespace": Namespace("http://posccaesar.org/imf/"),
+        "prefix": "pcaimf",
+        "namespace": Namespace("https://posccaesar.org/imf/"),
         "version": None,
         "created": None,
         "updated": None,
-        "title": "IMF - types",
+        "title": "IMF_types",
         "description": "IMF - types",
         "creator": None,
         "rights": None,

cognite/neat/rules/importers/_rdf/_imf2rules/_imf2properties.py

@@ -36,7 +36,7 @@ def parse_imf_to_properties(graph: Graph, language: str = "en") -> list[dict]:
     """
 
     query = """
-    SELECT ?class ?property ?name ?description ?valueType ?minCount ?maxCount ?default ?reference
+    SELECT DISTINCT ?class ?property ?name ?description ?valueType ?minCount ?maxCount ?default ?reference
     ?match ?comment ?propertyType
     WHERE
     {
@@ -48,7 +48,7 @@ def parse_imf_to_properties(graph: Graph, language: str = "en") -> list[dict]:
                 ?propertyShape sh:path ?imfProperty .
 
             OPTIONAL { ?imfProperty skos:prefLabel ?name . }
-            OPTIONAL { ?imfProperty skos:description ?description . }
+            OPTIONAL { ?imfProperty skos:definition | skos:description ?description . }
             OPTIONAL { ?imfProperty rdfs:range ?range . }
             OPTIONAL { ?imfProperty a ?type . }
             OPTIONAL { ?propertyShape sh:minCount ?minCardinality} .
@@ -62,22 +62,24 @@ def parse_imf_to_properties(graph: Graph, language: str = "en") -> list[dict]:
             ?imfClass a imf:AttributeType ;
                 ?imfProperty ?default .
 
-            # imf:predicate is required
-            BIND(IF(?imfProperty = <http://ns.imfid.org/imf#predicate>, 1, 0) AS ?minCardinality)
-
             # The following information is used to describe the attribute when it is connected to a block or a terminal
             # and not duplicated here.
-            FILTER(?imfProperty != rdf:type && ?imfProperty != skos:prefLabel && ?imfProperty != skos:description)
+            # Note: Bug in PCA has lead to the use non-existing term skos:description. This will be replaced
+            # with the correct skos:definition in the near future, so both terms are included here.
+            FILTER(?imfProperty != rdf:type && ?imfProperty != skos:prefLabel &&
+                ?imfProperty != skos:defintion && ?imfProperty != skos:description)
         }
 
         # Finding the last segment of the class IRI
         BIND(STR(?imfClass) AS ?classString)
-        BIND(REPLACE(?classString, "^.*[/#]([^/#]*)$", "$1") AS ?classSegment)
+        BIND(REPLACE(?classString, "^.*[/#]([^/#]*)$", "$1") AS ?tempClassSegment)
+        BIND(REPLACE(?tempClassSegment, "-", "_") AS ?classSegment)
         BIND(IF(CONTAINS(?classString, "imf/"), CONCAT("IMF_", ?classSegment) , ?classSegment) AS ?class)
 
         # Finding the last segment of the property IRI
         BIND(STR(?imfProperty) AS ?propertyString)
-        BIND(REPLACE(?propertyString, "^.*[/#]([^/#]*)$", "$1") AS ?propertySegment)
+        BIND(REPLACE(?propertyString, "^.*[/#]([^/#]*)$", "$1") AS ?tempPropertySegment)
+        BIND(REPLACE(?tempPropertySegment, "-", "_") AS ?propertySegment)
         BIND(IF(CONTAINS(?propertyString, "imf/"), CONCAT("IMF_", ?propertySegment) , ?propertySegment) AS ?property)
 
         # Set the value type for the property based on sh:class, sh:qualifiedValueType or rdfs:range
@@ -85,12 +87,17 @@ def parse_imf_to_properties(graph: Graph, language: str = "en") -> list[dict]:
 
         # Finding the last segment of value types
         BIND(STR(?valueIriType) AS ?valueTypeString)
-        BIND(REPLACE(?valueTypeString, "^.*[/#]([^/#]*)$", "$1") AS ?valueTypeSegment)
+        BIND(REPLACE(?valueTypeString, "^.*[/#]([^/#]*)$", "$1") AS ?tempValueTypeSegment)
+        BIND(REPLACE(?tempValueTypeSegment, "-", "_") AS ?valueTypeSegment)
         BIND(IF(CONTAINS(?valueTypeString, "imf/"), CONCAT("IMF_", ?valueTypeSegment) , ?valueTypeSegment)
             AS ?valueType)
 
-        # Helper variable to set property type if value type is not already set.
-        BIND(IF(BOUND(?type) && ?type = owl:DatatypeProperty, ?type , owl:ObjectProperty) AS ?propertyType)
+        # Helper variable to set owl datatype- or object-property if this is not already set.
+        BIND(IF( EXISTS {?imfProperty a ?tempPropertyType . FILTER(?tempPropertyType = owl:DatatypeProperty) },
+            owl:DatatypeProperty,
+            owl:ObjectProperty
+            )
+        AS ?propertyType)
 
         # Assert cardinality values if they do not exist
         BIND(IF(BOUND(?minCardinality), ?minCardinality, 0) AS ?minCount)

cognite/neat/rules/importers/_rdf/_imf2rules/_imf2rules.py

@@ -2,13 +2,14 @@
 there are loaders to TransformationRules pydantic class."""
 
 from pathlib import Path
-from typing import Literal, overload
 
 from rdflib import DC, DCTERMS, OWL, RDF, RDFS, SH, SKOS, Graph
 
 from cognite.neat.issues import IssueList
-from cognite.neat.rules.importers._base import BaseImporter, VerifiedRules
-from cognite.neat.rules.models import InformationRules, RoleTypes
+from cognite.neat.issues.errors import FileReadError
+from cognite.neat.rules._shared import ReadRules
+from cognite.neat.rules.importers._base import BaseImporter
+from cognite.neat.rules.models import InformationInputRules
 from cognite.neat.rules.models.data_types import _XSD_TYPES
 
 from ._imf2classes import parse_imf_to_classes
@@ -16,7 +17,7 @@ from ._imf2metadata import parse_imf_metadata
 from ._imf2properties import parse_imf_to_properties
 
 
-class IMFImporter(BaseImporter):
+class IMFImporter(BaseImporter[InformationInputRules]):
     """Convert SHACL shapes to tables/ transformation rules / Excel file.
 
         Args:
@@ -38,24 +39,14 @@ class IMFImporter(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:
+    ) -> ReadRules[InformationInputRules]:
         graph = Graph()
         try:
             graph.parse(self.filepath)
         except Exception as e:
-            raise Exception(f"Could not parse owl file: {e}") from e
+            return ReadRules(None, IssueList([FileReadError(self.filepath, f"Could not parse owl file: {e}")]), {})
 
         # bind key namespaces
         graph.bind("owl", OWL)
@@ -76,8 +67,8 @@ class IMFImporter(BaseImporter):
 
         components = make_components_compliant(components)
 
-        rules = InformationRules.model_validate(components)
-        return self._to_output(rules, IssueList(), errors, role)
+        rules = InformationInputRules.load(components)
+        return ReadRules(rules, IssueList(), {})
 
 
 def make_components_compliant(components: dict) -> dict:

cognite/neat/rules/importers/_rdf/_inference2rules.py

@@ -1,7 +1,7 @@
 from collections import Counter, defaultdict
 from datetime import datetime
 from pathlib import Path
-from typing import Literal, cast, overload
+from typing import cast
 
 from rdflib import Graph, Namespace, URIRef
 from rdflib import Literal as RdfLiteral
@@ -9,12 +9,12 @@ from rdflib import Literal as RdfLiteral
 from cognite.neat.constants import DEFAULT_NAMESPACE, get_default_prefixes
 from cognite.neat.issues import IssueList
 from cognite.neat.issues.errors import FileReadError
-from cognite.neat.rules.importers._base import BaseImporter, VerifiedRules, _handle_issues
-from cognite.neat.rules.models import InformationRules, RoleTypes
-from cognite.neat.rules.models._base import MatchType
+from cognite.neat.rules._shared import ReadRules
+from cognite.neat.rules.importers._base import BaseImporter
+from cognite.neat.rules.models._base_rules import MatchType
 from cognite.neat.rules.models.information import (
+    InformationInputRules,
     InformationMetadata,
-    InformationRulesInput,
 )
 from cognite.neat.store import NeatGraphStore
 from cognite.neat.utils.rdf_ import get_namespace, remove_namespace_from_uri, uri_to_short_form
@@ -44,7 +44,7 @@ INSTANCE_PROPERTIES_DEFINITION = """SELECT ?property (count(?property) as ?occur
                                     GROUP BY ?property ?dataType ?objectType"""
 
 
-class InferenceImporter(BaseImporter):
+class InferenceImporter(BaseImporter[InformationInputRules]):
     """Infers rules from a triple store.
 
     Rules inference through analysis of knowledge graph provided in various formats.
@@ -147,44 +147,21 @@ class InferenceImporter(BaseImporter):
     ) -> "InferenceImporter":
         raise NotImplementedError("JSON file format is not supported yet.")
 
-    @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:
+    ) -> ReadRules[InformationInputRules]:
         """
         Creates `Rules` object from the data for target role.
         """
 
         if self.issue_list.has_errors:
             # In case there were errors during the import, the to_rules method will return None
-            return self._return_or_raise(self.issue_list, errors)
+            return ReadRules(None, self.issue_list, {})
 
         rules_dict = self._to_rules_components()
 
-        with _handle_issues(self.issue_list) as future:
-            rules: InformationRules
-            rules = InformationRulesInput.load(rules_dict).as_rules()
-
-        if future.result == "failure" or self.issue_list.has_errors:
-            return self._return_or_raise(self.issue_list, errors)
-
-        return self._to_output(
-            rules,
-            self.issue_list,
-            errors=errors,
-            role=role,
-        )
+        rules = InformationInputRules.load(rules_dict)
+        return ReadRules(rules, self.issue_list, {})
 
     def _to_rules_components(
         self,

cognite/neat/rules/importers/_rdf/_owl2rules/_owl2rules.py

@@ -2,21 +2,22 @@
 there are loaders to TransformationRules pydantic class."""
 
 from pathlib import Path
-from typing import Literal, overload
 
 from rdflib import DC, DCTERMS, OWL, RDF, RDFS, SKOS, Graph
 
 from cognite.neat.issues import IssueList
-from cognite.neat.rules.importers._base import BaseImporter, VerifiedRules
+from cognite.neat.issues.errors import FileReadError
+from cognite.neat.rules._shared import ReadRules
+from cognite.neat.rules.importers._base import BaseImporter
 from cognite.neat.rules.importers._rdf._shared import make_components_compliant
-from cognite.neat.rules.models import InformationRules, RoleTypes
+from cognite.neat.rules.models import InformationInputRules
 
 from ._owl2classes import parse_owl_classes
 from ._owl2metadata import parse_owl_metadata
 from ._owl2properties import parse_owl_properties
 
 
-class OWLImporter(BaseImporter):
+class OWLImporter(BaseImporter[InformationInputRules]):
     """Convert OWL ontology to tables/ transformation rules / Excel file.
 
         Args:
@@ -37,24 +38,12 @@ class OWLImporter(BaseImporter):
     def __init__(self, filepath: Path):
         self.owl_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[InformationInputRules]:
         graph = Graph()
         try:
             graph.parse(self.owl_filepath)
         except Exception as e:
-            raise Exception(f"Could not parse owl file: {e}") from e
+            return ReadRules(None, IssueList([FileReadError(self.owl_filepath, f"Could not parse owl file: {e}")]), {})
 
         # bind key namespaces
         graph.bind("owl", OWL)
@@ -72,5 +61,5 @@ class OWLImporter(BaseImporter):
 
         components = make_components_compliant(components)
 
-        rules = InformationRules.model_validate(components)
-        return self._to_output(rules, IssueList(), errors, role)
+        rules = InformationInputRules.load(components)
+        return ReadRules(rules, IssueList(), {})

cognite/neat/rules/importers/_rdf/_shared.py

@@ -4,7 +4,7 @@ import numpy as np
 import pandas as pd
 from rdflib import OWL, Literal, Namespace
 
-from cognite.neat.rules.models._base import MatchType
+from cognite.neat.rules.models._base_rules import MatchType
 from cognite.neat.rules.models.data_types import _XSD_TYPES
 from cognite.neat.utils.rdf_ import remove_namespace_from_uri
 from cognite.neat.utils.regex_patterns import PATTERNS

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",