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

cognite/neat/rules/_shared.py

@@ -1,10 +1,61 @@
-from typing import TypeAlias
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Generic, TypeAlias, TypeVar
 
+from cognite.neat.issues import IssueList
 from cognite.neat.rules.models import (
     AssetRules,
     DMSRules,
     DomainRules,
     InformationRules,
 )
+from cognite.neat.rules.models.asset._rules_input import AssetInputRules
+from cognite.neat.rules.models.dms._rules_input import DMSInputRules
+from cognite.neat.rules.models.information._rules_input import InformationInputRules
 
-Rules: TypeAlias = DomainRules | InformationRules | DMSRules | AssetRules
+VerifiedRules: TypeAlias = DomainRules | InformationRules | DMSRules | AssetRules
+InputRules: TypeAlias = AssetInputRules | DMSInputRules | InformationInputRules
+Rules: TypeAlias = (
+    AssetInputRules | DMSInputRules | InformationInputRules | DomainRules | InformationRules | DMSRules | AssetRules
+)
+T_Rules = TypeVar("T_Rules", bound=Rules)
+T_VerifiedRules = TypeVar("T_VerifiedRules", bound=VerifiedRules)
+T_InputRules = TypeVar("T_InputRules", bound=InputRules)
+
+
+@dataclass
+class OutRules(Generic[T_Rules], ABC):
+    """This is a base class for all rule states."""
+
+    @abstractmethod
+    def get_rules(self) -> T_Rules | None:
+        """Get the rules from the state."""
+        raise NotImplementedError()
+
+
+@dataclass
+class JustRules(OutRules[T_Rules]):
+    """This represents a rule that exists"""
+
+    rules: T_Rules
+
+    def get_rules(self) -> T_Rules:
+        return self.rules
+
+
+@dataclass
+class MaybeRules(OutRules[T_Rules]):
+    """This represents a rule that may or may not exist"""
+
+    rules: T_Rules | None
+    issues: IssueList
+
+    def get_rules(self) -> T_Rules | None:
+        return self.rules
+
+
+@dataclass
+class ReadRules(MaybeRules[T_Rules]):
+    """This represents a rule that does not exist"""
+
+    read_context: dict[str, Any]

cognite/neat/rules/analysis/_base.py

@@ -9,7 +9,7 @@ from typing import Generic, TypeVar
 import pandas as pd
 from pydantic import BaseModel
 
-from cognite.neat.rules.models._base import BaseRules
+from cognite.neat.rules.models._base_rules import BaseRules
 from cognite.neat.rules.models._rdfpath import RDFPath
 from cognite.neat.rules.models.entities import (
     ClassEntity,

cognite/neat/rules/exporters/_base.py

@@ -5,47 +5,36 @@ from typing import Generic, TypeVar
 
 from cognite.client import CogniteClient
 
-from cognite.neat.rules._shared import Rules
-from cognite.neat.rules.models import DMSRules, InformationRules, RoleTypes
+from cognite.neat.rules._shared import T_VerifiedRules
 from cognite.neat.utils.auxiliary import class_html_doc
 from cognite.neat.utils.upload import UploadResult, UploadResultList
 
 T_Export = TypeVar("T_Export")
 
 
-class BaseExporter(ABC, Generic[T_Export]):
+class BaseExporter(ABC, Generic[T_VerifiedRules, T_Export]):
     _new_line = "\n"
     _encoding = "utf-8"
 
     @abstractmethod
-    def export_to_file(self, rules: Rules, filepath: Path) -> None:
+    def export_to_file(self, rules: T_VerifiedRules, filepath: Path) -> None:
         raise NotImplementedError
 
     @abstractmethod
-    def export(self, rules: Rules) -> T_Export:
+    def export(self, rules: T_VerifiedRules) -> T_Export:
         raise NotImplementedError
 
-    def _convert_to_output_role(self, rules: Rules, output_role: RoleTypes | None = None) -> Rules:
-        if rules.metadata.role is output_role or output_role is None:
-            return rules
-        elif output_role is RoleTypes.dms and isinstance(rules, InformationRules):
-            return rules.as_dms_rules()
-        elif output_role is RoleTypes.information and isinstance(rules, DMSRules):
-            return rules.as_information_rules()
-        else:
-            raise NotImplementedError(f"Role {output_role} is not supported for {type(rules).__name__} rules")
-
     @classmethod
     def _repr_html_(cls) -> str:
         return class_html_doc(cls, include_factory_methods=False)
 
 
-class CDFExporter(BaseExporter[T_Export]):
+class CDFExporter(BaseExporter[T_VerifiedRules, T_Export]):
     @abstractmethod
     def export_to_cdf_iterable(
-        self, rules: Rules, client: CogniteClient, dry_run: bool = False
+        self, rules: T_VerifiedRules, client: CogniteClient, dry_run: bool = False
     ) -> Iterable[UploadResult]:
         raise NotImplementedError
 
-    def export_to_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> UploadResultList:
+    def export_to_cdf(self, rules: T_VerifiedRules, client: CogniteClient, dry_run: bool = False) -> UploadResultList:
         return UploadResultList(self.export_to_cdf_iterable(rules, client, dry_run))

cognite/neat/rules/exporters/_rules2dms.py

@@ -16,10 +16,10 @@ from cognite.client.data_classes.data_modeling import (
 from cognite.client.exceptions import CogniteAPIError
 
 from cognite.neat.issues import IssueList
-from cognite.neat.issues.neat_warnings.resources import FailedLoadingResourcesWarning
-from cognite.neat.rules import issues
-from cognite.neat.rules._shared import Rules
-from cognite.neat.rules.models import InformationRules
+from cognite.neat.issues.warnings import (
+    PrincipleOneModelOneSpaceWarning,
+    ResourceRetrievalWarning,
+)
 from cognite.neat.rules.models.dms import DMSRules, DMSSchema, PipelineSchema
 from cognite.neat.utils.cdf.loaders import (
     ContainerLoader,
@@ -39,7 +39,7 @@ from ._base import CDFExporter
 Component: TypeAlias = Literal["all", "spaces", "data_models", "views", "containers", "node_types"]
 
 
-class DMSExporter(CDFExporter[DMSSchema]):
+class DMSExporter(CDFExporter[DMSRules, DMSSchema]):
     """Export rules to Cognite Data Fusion's Data Model Storage (DMS) service.
 
     Args:
@@ -80,7 +80,7 @@ class DMSExporter(CDFExporter[DMSSchema]):
         self.suppress_warnings = suppress_warnings
         self._schema: DMSSchema | None = None
 
-    def export_to_file(self, rules: Rules, filepath: Path) -> None:
+    def export_to_file(self, rules: DMSRules, filepath: Path) -> None:
         """Export the rules to a file(s).
 
         If the file is a directory, the components will be exported to separate files, otherwise they will be
@@ -95,12 +95,12 @@ class DMSExporter(CDFExporter[DMSSchema]):
         else:
             self._export_to_zip_file(filepath, rules)
 
-    def _export_to_directory(self, directory: Path, rules: Rules) -> None:
+    def _export_to_directory(self, directory: Path, rules: DMSRules) -> None:
         schema = self.export(rules)
         exclude = self._create_exclude_set()
         schema.to_directory(directory, exclude=exclude, new_line=self._new_line, encoding=self._encoding)
 
-    def _export_to_zip_file(self, filepath: Path, rules: Rules) -> None:
+    def _export_to_zip_file(self, filepath: Path, rules: DMSRules) -> None:
         if filepath.suffix not in {".zip"}:
             warnings.warn("File extension is not .zip, adding it to the file name", stacklevel=2)
             filepath = filepath.with_suffix(".zip")
@@ -115,16 +115,10 @@ class DMSExporter(CDFExporter[DMSSchema]):
             exclude = {"spaces", "data_models", "views", "containers", "node_types"} - self.export_components
         return exclude
 
-    def export(self, rules: Rules) -> DMSSchema:
-        if isinstance(rules, DMSRules):
-            dms_rules = rules
-        elif isinstance(rules, InformationRules):
-            dms_rules = rules.as_dms_rules()
-        else:
-            raise ValueError(f"{type(rules).__name__} cannot be exported to DMS")
-        return dms_rules.as_schema(include_pipeline=self.export_pipeline, instance_space=self.instance_space)
+    def export(self, rules: DMSRules) -> DMSSchema:
+        return rules.as_schema(include_pipeline=self.export_pipeline, instance_space=self.instance_space)
 
-    def delete_from_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadResult]:
+    def delete_from_cdf(self, rules: DMSRules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadResult]:
         to_export = self._prepare_exporters(rules, client)
 
         # we need to reverse order in which we are picking up the items to delete
@@ -168,7 +162,7 @@ class DMSExporter(CDFExporter[DMSSchema]):
             )
 
     def export_to_cdf_iterable(
-        self, rules: Rules, client: CogniteClient, dry_run: bool = False
+        self, rules: DMSRules, client: CogniteClient, dry_run: bool = False
     ) -> Iterable[UploadResult]:
         to_export = self._prepare_exporters(rules, client)
 
@@ -298,7 +292,10 @@ class DMSExporter(CDFExporter[DMSSchema]):
         if isinstance(loader, DataModelLoader):
             models = cast(list[DataModelApply], items)
             if other_models := self._exist_other_data_models(loader, models):
-                warning = issues.dms.OtherDataModelsInSpaceWarning(models[0].space, other_models)
+                warning = PrincipleOneModelOneSpaceWarning(
+                    f"There are multiple data models in the same space {models[0].space}. "
+                    f"Other data models in the space are {other_models}.",
+                )
                 if not self.suppress_warnings:
                     warnings.warn(warning, stacklevel=2)
                 issue_list.append(warning)
@@ -314,7 +311,7 @@ class DMSExporter(CDFExporter[DMSSchema]):
         try:
             data_models = loader.client.data_modeling.data_models.list(space=space, limit=25, all_versions=False)
         except CogniteAPIError as e:
-            warnings.warn(FailedLoadingResourcesWarning[str](frozenset({space}), "Space", str(e)), stacklevel=2)
+            warnings.warn(ResourceRetrievalWarning(frozenset({space}), "space", str(e)), stacklevel=2)
             return []
         else:
             return [

cognite/neat/rules/exporters/_rules2excel.py

@@ -12,11 +12,10 @@ from openpyxl.cell import MergedCell
 from openpyxl.styles import Alignment, Border, Font, PatternFill, Side
 from openpyxl.worksheet.worksheet import Worksheet
 
-from cognite.neat.rules._shared import Rules
+from cognite.neat.rules._shared import VerifiedRules
 from cognite.neat.rules.models import (
     DataModelType,
     ExtensionCategory,
-    RoleTypes,
     SchemaCompleteness,
     SheetEntity,
 )
@@ -27,7 +26,7 @@ from cognite.neat.rules.models.information import InformationMetadata
 from ._base import BaseExporter
 
 
-class ExcelExporter(BaseExporter[Workbook]):
+class ExcelExporter(BaseExporter[VerifiedRules, Workbook]):
     """Export rules to Excel.
 
     Args:
@@ -68,16 +67,14 @@ class ExcelExporter(BaseExporter[Workbook]):
         "Classes": "Definition of Classes",
         "Views": "Definition of Views",
         "Containers": "Definition of Containers",
+        "Nodes": "Definition of Nodes",
+        "Enum": "Definition of Enum Collections",
     }
     style_options = get_args(Style)
     dump_options = get_args(DumpOptions)
 
     def __init__(
-        self,
-        styling: Style = "default",
-        output_role: RoleTypes | None = None,
-        dump_as: DumpOptions = "user",
-        new_model_id: tuple[str, str] | None = None,
+        self, styling: Style = "default", dump_as: DumpOptions = "user", new_model_id: tuple[str, str] | None = None
     ):
         if styling not in self.style_options:
             raise ValueError(f"Invalid styling: {styling}. Valid options are {self.style_options}")
@@ -85,11 +82,10 @@ class ExcelExporter(BaseExporter[Workbook]):
             raise ValueError(f"Invalid dump_as: {dump_as}. Valid options are {self.dump_options}")
         self.styling = styling
         self._styling_level = self.style_options.index(styling)
-        self.output_role = output_role
         self.new_model_id = new_model_id
         self.dump_as = dump_as
 
-    def export_to_file(self, rules: Rules, filepath: Path) -> None:
+    def export_to_file(self, rules: VerifiedRules, filepath: Path) -> None:
         """Exports transformation rules to excel file."""
         data = self.export(rules)
         try:
@@ -98,8 +94,7 @@ class ExcelExporter(BaseExporter[Workbook]):
             data.close()
         return None
 
-    def export(self, rules: Rules) -> Workbook:
-        rules = self._convert_to_output_role(rules, self.output_role)
+    def export(self, rules: VerifiedRules) -> Workbook:
         workbook = Workbook()
         # Remove default sheet named "Sheet"
         workbook.remove(workbook["Sheet"])
@@ -147,7 +142,7 @@ class ExcelExporter(BaseExporter[Workbook]):
         self,
         workbook: Workbook,
         dumped_rules: dict[str, Any],
-        rules: Rules,
+        rules: VerifiedRules,
         sheet_prefix: str = "",
     ):
         for sheet_name, headers in rules.headers_by_sheet(by_alias=True).items():
@@ -279,9 +274,7 @@ class _MetadataCreator:
 
         new_metadata = self._create_new_info(now)
         if isinstance(metadata, DMSMetadata):
-            from cognite.neat.rules.models.information._converter import (
-                _InformationRulesConverter,
-            )
+            from cognite.neat.rules.transformers._converters import _InformationRulesConverter
 
             output_metadata: DMSMetadata | InformationMetadata = _InformationRulesConverter._convert_metadata_to_dms(
                 new_metadata

cognite/neat/rules/exporters/_rules2ontology.py

@@ -9,20 +9,12 @@ from rdflib import DCTERMS, OWL, RDF, RDFS, XSD, BNode, Graph, Literal, Namespac
 from rdflib.collection import Collection as GraphCollection
 
 from cognite.neat.constants import DEFAULT_NAMESPACE as NEAT_NAMESPACE
-from cognite.neat.rules.analysis import InformationAnalysis
-from cognite.neat.rules.issues.ontology import (
-    MetadataSheetNamespaceNotDefinedError,
-    MissingDataModelPrefixOrNamespaceWarning,
-    OntologyMultiDefinitionPropertyWarning,
-    OntologyMultiDomainPropertyWarning,
-    OntologyMultiLabeledPropertyWarning,
-    OntologyMultiRangePropertyWarning,
-    OntologyMultiTypePropertyWarning,
-    PrefixMissingError,
-    PropertiesDefinedMultipleTimesError,
-    PropertyDefinitionsNotForSamePropertyError,
+from cognite.neat.issues import MultiValueError
+from cognite.neat.issues.errors import (
+    PropertyDefinitionDuplicatedError,
 )
-from cognite.neat.rules.models import DMSRules
+from cognite.neat.issues.warnings import PropertyDefinitionDuplicatedWarning
+from cognite.neat.rules.analysis import InformationAnalysis
 from cognite.neat.rules.models.data_types import DataType
 from cognite.neat.rules.models.entities import ClassEntity, EntityTypes
 from cognite.neat.rules.models.information import (
@@ -31,43 +23,40 @@ from cognite.neat.rules.models.information import (
     InformationProperty,
     InformationRules,
 )
-from cognite.neat.utils.auxiliary import generate_exception_report
 from cognite.neat.utils.rdf_ import remove_namespace_from_uri
 
 from ._base import BaseExporter
-from ._validation import are_properties_redefined
+from ._validation import duplicated_properties
 
 if sys.version_info >= (3, 11):
     from typing import Self
 else:
     from typing_extensions import Self
 
-from cognite.neat.rules._shared import Rules
 
-
-class GraphExporter(BaseExporter[Graph], ABC):
-    def export_to_file(self, rules: Rules, filepath: Path) -> None:
+class GraphExporter(BaseExporter[InformationRules, Graph], ABC):
+    def export_to_file(self, rules: InformationRules, filepath: Path) -> None:
         self.export(rules).serialize(destination=filepath, encoding=self._encoding, newline=self._new_line)
 
 
 class OWLExporter(GraphExporter):
-    """Exports rules to an OWL ontology."""
+    """Exports verified information rules to an OWL ontology."""
 
-    def export(self, rules: Rules) -> Graph:
+    def export(self, rules: InformationRules) -> Graph:
         return Ontology.from_rules(rules).as_owl()
 
 
 class SHACLExporter(GraphExporter):
     """Exports rules to a SHACL graph."""
 
-    def export(self, rules: Rules) -> Graph:
+    def export(self, rules: InformationRules) -> Graph:
         return Ontology.from_rules(rules).as_shacl()
 
 
 class SemanticDataModelExporter(GraphExporter):
-    """Exports rules to a semantic data model."""
+    """Exports verified information rules to a semantic data model."""
 
-    def export(self, rules: Rules) -> Graph:
+    def export(self, rules: InformationRules) -> Graph:
         return Ontology.from_rules(rules).as_semantic_data_model()
 
 
@@ -94,34 +83,30 @@ class Ontology(OntologyModel):
     prefixes: dict[str, Namespace]
 
     @classmethod
-    def from_rules(cls, input_rules: Rules) -> Self:
+    def from_rules(cls, rules: InformationRules) -> Self:
         """
         Generates an ontology from a set of transformation rules.
 
         Args:
-            input_rules: The rules to generate the ontology from.
+            rules: The rules to generate the ontology from.
 
         Returns:
             An instance of Ontology.
         """
-        if isinstance(input_rules, InformationRules):
-            rules = input_rules
-        elif isinstance(input_rules, DMSRules):
-            rules = input_rules.as_information_rules()
-        else:
-            raise ValueError(f"{type(input_rules).__name__} cannot be exported to Ontology")
-
-        properties_redefined, redefinition_warnings = are_properties_redefined(rules, return_report=True)
-        if properties_redefined:
-            raise PropertiesDefinedMultipleTimesError(
-                report=generate_exception_report(redefinition_warnings)
-            ).as_exception()
-
-        if rules.prefixes is None:
-            raise PrefixMissingError().as_exception()
-
-        if rules.metadata.namespace is None:
-            raise MissingDataModelPrefixOrNamespaceWarning()
+        if duplicates := duplicated_properties(rules.properties):
+            errors = []
+            for (class_, property_), definitions in duplicates.items():
+                errors.append(
+                    PropertyDefinitionDuplicatedError(
+                        class_,
+                        "class",
+                        property_,
+                        frozenset({str(definition[1].value_type) for definition in definitions}),
+                        tuple(definition[0] for definition in definitions),
+                        "rows",
+                    )
+                )
+            raise MultiValueError(errors)
 
         class_dict = InformationAnalysis(rules).as_class_dict()
         return cls(
@@ -184,9 +169,6 @@ class Ontology(OntologyModel):
         for prefix, namespace in self.prefixes.items():
             owl.bind(prefix, namespace)
 
-        if self.metadata.namespace is None:
-            raise MetadataSheetNamespaceNotDefinedError().as_exception()
-
         owl.add((URIRef(self.metadata.namespace), RDF.type, OWL.Ontology))
         for property_ in self.properties:
             for triple in property_.triples:
@@ -233,8 +215,6 @@ class OWLMetadata(InformationMetadata):
     @property
     def triples(self) -> list[tuple]:
         # Mandatory triples originating from Metadata mandatory fields
-        if self.namespace is None:
-            raise MetadataSheetNamespaceNotDefinedError().as_exception()
         triples: list[tuple] = [
             (URIRef(self.namespace), DCTERMS.hasVersion, Literal(self.version)),
             (URIRef(self.namespace), OWL.versionInfo, Literal(self.version)),
@@ -323,16 +303,17 @@ class OWLProperty(OntologyModel):
     range_: set[URIRef]
     namespace: Namespace
 
-    @staticmethod
-    def same_property_id(definitions: list[InformationProperty]) -> bool:
-        return len({definition.property_ for definition in definitions}) == 1
-
     @classmethod
     def from_list_of_properties(cls, definitions: list[InformationProperty], namespace: Namespace) -> "OWLProperty":
         """Here list of properties is a list of properties with the same id, but different definitions."""
-
-        if not cls.same_property_id(definitions):
-            raise PropertyDefinitionsNotForSamePropertyError().as_exception()
+        property_ids = {definition.property_ for definition in definitions}
+        if len(property_ids) != 1:
+            raise PropertyDefinitionDuplicatedError(
+                definitions[0].class_,
+                "class",
+                definitions[0].property_,
+                frozenset(property_ids),
+            )
 
         owl_property = cls.model_construct(
             id_=namespace[definitions[0].property_],
@@ -363,8 +344,15 @@ class OWLProperty(OntologyModel):
     def is_multi_type(cls, v, info: ValidationInfo):
         if len(v) > 1:
             warnings.warn(
-                OntologyMultiTypePropertyWarning(
-                    remove_namespace_from_uri(info.data["id_"]), [remove_namespace_from_uri(t) for t in v]
+                PropertyDefinitionDuplicatedWarning(
+                    remove_namespace_from_uri(info.data["id"]),
+                    "class",
+                    "type",
+                    frozenset({remove_namespace_from_uri(t) for t in v}),
+                    "This warning occurs when a same property is define for two object/classes where"
+                    " its expected value type is different in one definition, e.g. acts as an edge, while in "
+                    "other definition acts as and attribute",
+                    "If a property takes different value types for different objects, simply define new property",
                 ),
                 stacklevel=2,
             )
@@ -374,8 +362,14 @@ class OWLProperty(OntologyModel):
     def is_multi_range(cls, v, info: ValidationInfo):
         if len(v) > 1:
             warnings.warn(
-                OntologyMultiRangePropertyWarning(
-                    remove_namespace_from_uri(info.data["id_"]), [remove_namespace_from_uri(t) for t in v]
+                PropertyDefinitionDuplicatedWarning(
+                    remove_namespace_from_uri(info.data["id_"]),
+                    "class",
+                    "range",
+                    frozenset({remove_namespace_from_uri(t) for t in v}),
+                    "This warning occurs when a property takes range of "
+                    "values which consists of union of multiple value types.",
+                    "If value types for different objects, simply define new property",
                 ),
                 stacklevel=2,
             )
@@ -385,8 +379,15 @@ class OWLProperty(OntologyModel):
     def is_multi_domain(cls, v, info: ValidationInfo):
         if len(v) > 1:
             warnings.warn(
-                OntologyMultiDomainPropertyWarning(
-                    remove_namespace_from_uri(info.data["id_"]), [remove_namespace_from_uri(t) for t in v]
+                PropertyDefinitionDuplicatedWarning(
+                    remove_namespace_from_uri(info.data["id_"]),
+                    "class",
+                    "domain",
+                    frozenset({remove_namespace_from_uri(t) for t in v}),
+                    "This warning occurs when a same property is define for two object/classes where"
+                    " its expected value type is different in one definition, e.g. acts as an edge, while in "
+                    "other definition acts as and attribute",
+                    "If value types for different objects, simply define new property",
                 ),
                 stacklevel=2,
             )
@@ -396,7 +397,13 @@ class OWLProperty(OntologyModel):
     def has_multi_name(cls, v, info: ValidationInfo):
         if len(v) > 1:
             warnings.warn(
-                OntologyMultiLabeledPropertyWarning(remove_namespace_from_uri(info.data["id_"]), v),
+                PropertyDefinitionDuplicatedWarning(
+                    remove_namespace_from_uri(info.data["id_"]),
+                    "class",
+                    "label",
+                    frozenset(v),
+                    f"Only the first label (name) will be used, {v[0]}",
+                ),
                 stacklevel=2,
             )
         return v
@@ -405,7 +412,13 @@ class OWLProperty(OntologyModel):
     def has_multi_comment(cls, v, info: ValidationInfo):
         if len(v) > 1:
             warnings.warn(
-                OntologyMultiDefinitionPropertyWarning(remove_namespace_from_uri(info.data["id_"])),
+                PropertyDefinitionDuplicatedWarning(
+                    remove_namespace_from_uri(info.data["id_"]),
+                    "class",
+                    "comment",
+                    frozenset(v),
+                    "All definitions will be concatenated to form a single definition.",
+                ),
                 stacklevel=2,
             )
         return v

cognite/neat/rules/exporters/_rules2yaml.py

@@ -5,14 +5,13 @@ from typing import Literal, get_args
 
 import yaml
 
-from cognite.neat.rules._shared import Rules
-from cognite.neat.rules.models import RoleTypes
+from cognite.neat.rules._shared import VerifiedRules
 
 from ._base import BaseExporter
 
 
-class YAMLExporter(BaseExporter[str]):
-    """Export rules to YAML.
+class YAMLExporter(BaseExporter[VerifiedRules, str]):
+    """Export rules (Information, DMS or Domain) to YAML.
 
     Args:
         files: The number of files to output. Defaults to "single".
@@ -38,16 +37,15 @@ class YAMLExporter(BaseExporter[str]):
     file_option = get_args(Files)
     format_option = get_args(Format)
 
-    def __init__(self, files: Files = "single", output: Format = "yaml", output_role: RoleTypes | None = None):
+    def __init__(self, files: Files = "single", output: Format = "yaml"):
         if files not in self.file_option:
             raise ValueError(f"Invalid files: {files}. Valid options are {self.file_option}")
         if output not in self.format_option:
             raise ValueError(f"Invalid output: {output}. Valid options are {self.format_option}")
         self.files = files
         self.output = output
-        self.output_role = output_role
 
-    def export_to_file(self, rules: Rules, filepath: Path) -> None:
+    def export_to_file(self, rules: VerifiedRules, filepath: Path) -> None:
         """Exports transformation rules to YAML/JSON file(s)."""
         if self.files == "single":
             if filepath.suffix != f".{self.output}":
@@ -57,7 +55,7 @@ class YAMLExporter(BaseExporter[str]):
         else:
             raise NotImplementedError(f"Exporting to {self.files} files is not supported")
 
-    def export(self, rules: Rules) -> str:
+    def export(self, rules: VerifiedRules) -> str:
         """Export rules to YAML (or JSON) format.
 
         Args:
@@ -66,7 +64,6 @@ class YAMLExporter(BaseExporter[str]):
         Returns:
             str: The rules in YAML (or JSON) format.
         """
-        rules = self._convert_to_output_role(rules, self.output_role)
         # model_dump_json ensures that the output is in JSON format,
         # if we don't do this, we will get Enums and other types that are not serializable to YAML
         json_output = rules.dump(mode="json")