cognite-neat 0.99.0__py3-none-any.whl → 0.100.0__py3-none-any.whl

cognite/neat/_rules/transformers/_converters.py

@@ -4,7 +4,7 @@ from abc import ABC, abstractmethod
 from collections import Counter, defaultdict
 from collections.abc import Collection, Mapping
 from datetime import date, datetime
-from typing import Literal, TypeVar, cast, overload
+from typing import ClassVar, Literal, TypeVar, cast, overload
 
 from cognite.client.data_classes import data_modeling as dms
 from cognite.client.data_classes.data_modeling import DataModelId, DataModelIdentifier, ViewId
@@ -15,6 +15,7 @@ from cognite.neat._constants import (
 )
 from cognite.neat._issues._base import IssueList
 from cognite.neat._issues.errors import NeatValueError
+from cognite.neat._issues.warnings import NeatValueWarning
 from cognite.neat._issues.warnings._models import (
     EnterpriseModelNotBuildOnTopOfCDMWarning,
     SolutionModelBuildOnTopOfCDMWarning,
@@ -34,7 +35,7 @@ from cognite.neat._rules.models import (
     SheetList,
     data_types,
 )
-from cognite.neat._rules.models.data_types import DataType, String
+from cognite.neat._rules.models.data_types import AnyURI, DataType, String
 from cognite.neat._rules.models.dms import DMSMetadata, DMSProperty, DMSView
 from cognite.neat._rules.models.dms._rules import DMSContainer
 from cognite.neat._rules.models.entities import (
@@ -249,11 +250,12 @@ class PrefixEntities(RulesTransformer[InputRules, InputRules]):  # type: ignore[
 class InformationToDMS(ConversionTransformer[InformationRules, DMSRules]):
     """Converts InformationRules to DMSRules."""
 
-    def __init__(self, ignore_undefined_value_types: bool = False):
+    def __init__(self, ignore_undefined_value_types: bool = False, mode: Literal["edge_properties"] | None = None):
         self.ignore_undefined_value_types = ignore_undefined_value_types
+        self.mode = mode
 
     def _transform(self, rules: InformationRules) -> DMSRules:
-        return _InformationRulesConverter(rules).as_dms_rules(self.ignore_undefined_value_types)
+        return _InformationRulesConverter(rules).as_dms_rules(self.ignore_undefined_value_types, self.mode)
 
 
 class DMSToInformation(ConversionTransformer[DMSRules, InformationRules]):
@@ -637,11 +639,15 @@ class ReduceCogniteModel(RulesTransformer[DMSRules, DMSRules]):
 
 
 class _InformationRulesConverter:
+    _edge_properties: ClassVar[frozenset[str]] = frozenset({"endNode", "end_node", "startNode", "start_node"})
+
     def __init__(self, information: InformationRules):
         self.rules = information
         self.property_count_by_container: dict[ContainerEntity, int] = defaultdict(int)
 
-    def as_dms_rules(self, ignore_undefined_value_types: bool = False) -> "DMSRules":
+    def as_dms_rules(
+        self, ignore_undefined_value_types: bool = False, mode: Literal["edge_properties"] | None = None
+    ) -> "DMSRules":
         from cognite.neat._rules.models.dms._rules import (
             DMSContainer,
             DMSProperty,
@@ -652,27 +658,55 @@ class _InformationRulesConverter:
         info_metadata = self.rules.metadata
         default_version = info_metadata.version
         default_space = self._to_space(info_metadata.prefix)
-        metadata = self._convert_metadata_to_dms(info_metadata)
+        dms_metadata = self._convert_metadata_to_dms(info_metadata)
+        edge_classes: set[ClassEntity] = set()
+        property_to_edge: dict[tuple[ClassEntity, str], ClassEntity] = {}
+        end_node_by_edge: dict[ClassEntity, ClassEntity] = {}
+        if mode == "edge_properties":
+            edge_classes = {
+                cls_.class_ for cls_ in self.rules.classes if cls_.implements and cls_.implements[0].suffix == "Edge"
+            }
+            property_to_edge = {
+                (prop.class_, prop.property_): prop.value_type
+                for prop in self.rules.properties
+                if prop.value_type in edge_classes and isinstance(prop.value_type, ClassEntity)
+            }
+            end_node_by_edge = {
+                prop.class_: prop.value_type
+                for prop in self.rules.properties
+                if prop.class_ in edge_classes
+                and (prop.property_ == "endNode" or prop.property_ == "end_node")
+                and isinstance(prop.value_type, ClassEntity)
+            }
 
         properties_by_class: dict[ClassEntity, list[DMSProperty]] = defaultdict(list)
         referenced_containers: dict[ContainerEntity, Counter[ClassEntity]] = defaultdict(Counter)
         for prop in self.rules.properties:
             if ignore_undefined_value_types and isinstance(prop.value_type, UnknownEntity):
                 continue
-            dms_property = self._as_dms_property(prop, default_space, default_version)
+            if prop.class_ in edge_classes and prop.property_ in self._edge_properties:
+                continue
+            dms_property = self._as_dms_property(
+                prop, default_space, default_version, edge_classes, property_to_edge, end_node_by_edge
+            )
             properties_by_class[prop.class_].append(dms_property)
             if dms_property.container:
                 referenced_containers[dms_property.container][prop.class_] += 1
 
-        views: list[DMSView] = [
-            DMSView(
+        views: list[DMSView] = []
+
+        for cls_ in self.rules.classes:
+            dms_view = DMSView(
                 name=cls_.name,
                 view=cls_.class_.as_view_entity(default_space, default_version),
                 description=cls_.description,
-                implements=self._get_view_implements(cls_, info_metadata),
+                implements=self._get_view_implements(cls_, info_metadata, mode),
             )
-            for cls_ in self.rules.classes
-        ]
+
+            dms_view.logical = cls_.neatId
+            cls_.physical = dms_view.neatId
+
+            views.append(dms_view)
 
         class_by_entity = {cls_.class_: cls_ for cls_ in self.rules.classes}
 
@@ -696,7 +730,7 @@ class _InformationRulesConverter:
             containers.append(container)
 
         return DMSRules(
-            metadata=metadata,
+            metadata=dms_metadata,
             properties=SheetList[DMSProperty]([prop for prop_set in properties_by_class.values() for prop in prop_set]),
             views=SheetList[DMSView](views),
             containers=SheetList[DMSContainer](containers),
@@ -724,7 +758,7 @@ class _InformationRulesConverter:
             DMSMetadata,
         )
 
-        return DMSMetadata(
+        dms_metadata = DMSMetadata(
             space=metadata.space,
             version=metadata.version,
             external_id=metadata.external_id,
@@ -734,74 +768,138 @@ class _InformationRulesConverter:
             updated=metadata.updated,
         )
 
-    def _as_dms_property(self, prop: InformationProperty, default_space: str, default_version: str) -> "DMSProperty":
-        """This creates the first"""
+        dms_metadata.logical = metadata.identifier
+        metadata.physical = dms_metadata.identifier
+
+        return dms_metadata
 
+    def _as_dms_property(
+        self,
+        info_property: InformationProperty,
+        default_space: str,
+        default_version: str,
+        edge_classes: set[ClassEntity],
+        property_to_edge: dict[tuple[ClassEntity, str], ClassEntity],
+        end_node_by_edge: dict[ClassEntity, ClassEntity],
+    ) -> "DMSProperty":
         from cognite.neat._rules.models.dms._rules import DMSProperty
 
         # returns property type, which can be ObjectProperty or DatatypeProperty
-        value_type: DataType | ViewEntity | DMSUnknownEntity
+        value_type = self._get_value_type(
+            info_property,
+            default_space,
+            default_version,
+            edge_classes,
+            end_node_by_edge,
+        )
+
+        connection = self._get_connection(info_property, value_type, property_to_edge, default_space, default_version)
+
+        container: ContainerEntity | None = None
+        container_property: str | None = None
+        is_list: bool | None = info_property.is_list
+        nullable: bool | None = not info_property.is_mandatory
+        if isinstance(connection, EdgeEntity):
+            nullable = None
+        elif connection == "direct":
+            nullable = True
+            container, container_property = self._get_container(info_property, default_space)
+        else:
+            container, container_property = self._get_container(info_property, default_space)
+
+        dms_property = DMSProperty(
+            name=info_property.name,
+            value_type=value_type,
+            nullable=nullable,
+            is_list=is_list,
+            connection=connection,
+            default=info_property.default,
+            container=container,
+            container_property=container_property,
+            view=info_property.class_.as_view_entity(default_space, default_version),
+            view_property=info_property.property_,
+        )
+
+        # linking
+        dms_property.logical = info_property.neatId
+        info_property.physical = dms_property.neatId
+
+        return dms_property
+
+    @staticmethod
+    def _get_connection(
+        prop: InformationProperty,
+        value_type: DataType | ViewEntity | DMSUnknownEntity,
+        property_to_edge: dict[tuple[ClassEntity, str], ClassEntity],
+        default_space: str,
+        default_version: str,
+    ) -> Literal["direct"] | ReverseConnectionEntity | EdgeEntity | None:
+        if isinstance(value_type, ViewEntity) and (prop.class_, prop.property_) in property_to_edge:
+            edge_properties = property_to_edge[(prop.class_, prop.property_)]
+            return EdgeEntity(properties=edge_properties.as_view_entity(default_space, default_version))
+        if isinstance(value_type, ViewEntity) and prop.is_list:
+            return EdgeEntity()
+        elif isinstance(value_type, ViewEntity):
+            return "direct"
+        # defaulting to direct connection
+        elif isinstance(value_type, DMSUnknownEntity):
+            return "direct"
+        return None
+
+    def _get_value_type(
+        self,
+        prop: InformationProperty,
+        default_space: str,
+        default_version: str,
+        edge_classes: set[ClassEntity],
+        end_node_by_edge: dict[ClassEntity, ClassEntity],
+    ) -> DataType | ViewEntity | DMSUnknownEntity:
         if isinstance(prop.value_type, DataType):
-            value_type = prop.value_type
+            return prop.value_type
 
         # UnknownEntity should  resolve to DMSUnknownEntity
         # meaning end node type is unknown
         elif isinstance(prop.value_type, UnknownEntity):
-            value_type = DMSUnknownEntity()
-
+            return DMSUnknownEntity()
+
+        elif isinstance(prop.value_type, ClassEntity) and (prop.value_type in edge_classes):
+            if prop.value_type in end_node_by_edge:
+                return end_node_by_edge[prop.value_type].as_view_entity(default_space, default_version)
+            warnings.warn(
+                NeatValueWarning(
+                    f"Edge class {prop.value_type} does not have 'endNode' property, defaulting to DMSUnknownEntity"
+                ),
+                stacklevel=2,
+            )
+            return DMSUnknownEntity()
         elif isinstance(prop.value_type, ClassEntity):
-            value_type = prop.value_type.as_view_entity(default_space, default_version)
+            return prop.value_type.as_view_entity(default_space, default_version)
 
         elif isinstance(prop.value_type, MultiValueTypeInfo):
             # Multi Object type should resolve to DMSUnknownEntity
             # meaning end node type is unknown
             if prop.value_type.is_multi_object_type():
-                value_type = DMSUnknownEntity()
+                non_unknown = [type_ for type_ in prop.value_type.types if isinstance(type_, UnknownEntity)]
+                if list(non_unknown) == 1:
+                    #
+                    return non_unknown[0].as_view_entity(default_space, default_version)
+                return DMSUnknownEntity()
 
             # Multi Data type should resolve to a single data type, or it should
             elif prop.value_type.is_multi_data_type():
-                value_type = self.convert_multi_data_type(prop.value_type)
+                return self.convert_multi_data_type(prop.value_type)
 
             # Mixed types default to string
             else:
-                value_type = String()
-
-        else:
-            raise ValueError(f"Unsupported value type: {prop.value_type.type_}")
+                non_any_uri = [type_ for type_ in prop.value_type.types if type_ != AnyURI()]
+                if list(non_any_uri) == 1:
+                    if isinstance(non_any_uri[0], ClassEntity):
+                        return non_any_uri[0].as_view_entity(default_space, default_version)
+                    else:
+                        return non_any_uri[0]
+                return String()
 
-        connection: Literal["direct"] | ReverseConnectionEntity | EdgeEntity | None = None
-        if isinstance(value_type, ViewEntity):
-            # Default connection type.
-            connection = EdgeEntity() if prop.is_list else "direct"
-
-        # defaulting to direct connection
-        elif isinstance(value_type, DMSUnknownEntity):
-            connection = "direct"
-
-        container: ContainerEntity | None = None
-        container_property: str | None = None
-        is_list: bool | None = prop.is_list
-        nullable: bool | None = not prop.is_mandatory
-        if isinstance(connection, EdgeEntity):
-            nullable = None
-        elif connection == "direct":
-            nullable = True
-            container, container_property = self._get_container(prop, default_space)
-        else:
-            container, container_property = self._get_container(prop, default_space)
-
-        return DMSProperty(
-            name=prop.name,
-            value_type=value_type,
-            nullable=nullable,
-            is_list=is_list,
-            connection=connection,
-            default=prop.default,
-            container=container,
-            container_property=container_property,
-            view=prop.class_.as_view_entity(default_space, default_version),
-            view_property=prop.property_,
-        )
+        raise ValueError(f"Unsupported value type: {prop.value_type.type_}")
 
     @classmethod
     def _to_space(cls, prefix: str) -> str:
@@ -823,9 +921,13 @@ class _InformationRulesConverter:
         self.property_count_by_container[container_entity] += 1
         return container_entity, prop.property_
 
-    def _get_view_implements(self, cls_: InformationClass, metadata: InformationMetadata) -> list[ViewEntity]:
+    def _get_view_implements(
+        self, cls_: InformationClass, metadata: InformationMetadata, mode: Literal["edge_properties"] | None
+    ) -> list[ViewEntity]:
         implements = []
         for parent in cls_.implements or []:
+            if mode == "edge_properties" and parent.suffix == "Edge":
+                continue
             view_entity = parent.as_view_entity(metadata.prefix, metadata.version)
             implements.append(view_entity)
         return implements

cognite/neat/_rules/transformers/_mapping.py

@@ -4,13 +4,16 @@ from collections import defaultdict
 from functools import cached_property
 from typing import Any, ClassVar, Literal
 
-from cognite.neat._issues.errors import NeatValueError
+from cognite.client import data_modeling as dm
+
+from cognite.neat._client import NeatClient
+from cognite.neat._issues.errors import CDFMissingClientError, NeatValueError, ResourceNotFoundError
 from cognite.neat._issues.warnings import NeatValueWarning, PropertyOverwritingWarning
 from cognite.neat._rules._shared import JustRules, OutRules
 from cognite.neat._rules.models import DMSRules, SheetList
 from cognite.neat._rules.models.data_types import Enum
 from cognite.neat._rules.models.dms import DMSEnum, DMSProperty, DMSView
-from cognite.neat._rules.models.entities import ViewEntity
+from cognite.neat._rules.models.entities import ContainerEntity, ViewEntity
 
 from ._base import RulesTransformer
 
@@ -208,3 +211,130 @@ class RuleMapper(RulesTransformer[DMSRules, DMSRules]):
                     # These are used for warnings so we use the alias to make it more readable for the user
                     conflicts.append(mapping_prop.model_fields[field_name].alias or field_name)
         return to_overwrite, conflicts
+
+
+class AsParentPropertyId(RulesTransformer[DMSRules, DMSRules]):
+    """Looks up all view properties that map to the same container property,
+    and changes the child view property id to match the parent property id.
+    """
+
+    def __init__(self, client: NeatClient | None = None) -> None:
+        self._client = client
+
+    def transform(self, rules: DMSRules | OutRules[DMSRules]) -> JustRules[DMSRules]:
+        input_rules = self._to_rules(rules)
+        new_rules = input_rules.model_copy(deep=True)
+        new_rules.metadata.version += "_as_parent_name"
+
+        path_by_view = self._inheritance_path_by_view(new_rules)
+        view_by_container_property = self._view_by_container_properties(new_rules)
+
+        parent_view_property_by_container_property = self._get_parent_view_property_by_container_property(
+            path_by_view, view_by_container_property
+        )
+
+        for prop in new_rules.properties:
+            if prop.container and prop.container_property:
+                if parent_name := parent_view_property_by_container_property.get(
+                    (prop.container, prop.container_property)
+                ):
+                    prop.view_property = parent_name
+
+        return JustRules(new_rules)
+
+    # Todo: Move into Probe class. Note this means that the Probe class must take a NeatClient as an argument.
+    def _inheritance_path_by_view(self, rules: DMSRules) -> dict[ViewEntity, list[ViewEntity]]:
+        parents_by_view: dict[ViewEntity, list[ViewEntity]] = {view.view: view.implements or [] for view in rules.views}
+
+        path_by_view: dict[ViewEntity, list[ViewEntity]] = {}
+        for view in rules.views:
+            path_by_view[view.view] = self._get_inheritance_path(
+                view.view, parents_by_view, rules.metadata.as_data_model_id()
+            )
+        return path_by_view
+
+    def _get_inheritance_path(
+        self, view: ViewEntity, parents_by_view: dict[ViewEntity, list[ViewEntity]], data_model_id: dm.DataModelId
+    ) -> list[ViewEntity]:
+        if parents_by_view.get(view) == []:
+            # We found the root.
+            return [view]
+        if view not in parents_by_view and self._client is not None:
+            # Lookup the parent
+            view_id = view.as_id()
+            read_views = self._client.loaders.views.retrieve([view_id])
+            if not read_views:
+                # Warning? Should be caught by validation
+                raise ResourceNotFoundError(view_id, "view", data_model_id, "data model")
+            parent_view_latest = max(read_views, key=lambda view: view.created_time)
+            parents_by_view[ViewEntity.from_id(parent_view_latest.as_id())] = [
+                ViewEntity.from_id(grand_parent) for grand_parent in parent_view_latest.implements or []
+            ]
+        elif view not in parents_by_view:
+            raise CDFMissingClientError(
+                f"The data model {data_model_id} is referencing a view that is not in the data model."
+                f"Please provide a client to lookup the view."
+            )
+
+        inheritance_path = [view]
+        seen = {view}
+        if view in parents_by_view:
+            for parent in parents_by_view[view]:
+                parent_path = self._get_inheritance_path(parent, parents_by_view, data_model_id)
+                inheritance_path.extend([p for p in parent_path if p not in seen])
+                seen.update(parent_path)
+        return inheritance_path
+
+    def _view_by_container_properties(
+        self, rules: DMSRules
+    ) -> dict[tuple[ContainerEntity, str], list[tuple[ViewEntity, str]]]:
+        view_properties_by_container_properties: dict[tuple[ContainerEntity, str], list[tuple[ViewEntity, str]]] = (
+            defaultdict(list)
+        )
+        view_with_properties: set[ViewEntity] = set()
+        for prop in rules.properties:
+            if not prop.container or not prop.container_property:
+                continue
+            view_properties_by_container_properties[(prop.container, prop.container_property)].append(
+                (prop.view, prop.view_property)
+            )
+            view_with_properties.add(prop.view)
+
+        # We need to look up all parent properties.
+        to_lookup = {view.view.as_id() for view in rules.views if view.view not in view_with_properties}
+        if to_lookup and self._client is None:
+            raise CDFMissingClientError(
+                f"Views {to_lookup} are not in the data model. Please provide a client to lookup the views."
+            )
+        elif to_lookup and self._client:
+            read_views = self._client.loaders.views.retrieve(list(to_lookup), include_ancestor=True)
+            write_views = [self._client.loaders.views.as_write(read_view) for read_view in read_views]
+            # We use the write/request format of the views as the read/response format contains all properties
+            # including ancestor properties. The goal is to find the property name used in the parent
+            # and thus we cannot have that repeated in the child views.
+            for write_view in write_views:
+                view_id = write_view.as_id()
+                view_entity = ViewEntity.from_id(view_id)
+
+                for property_id, property_ in (write_view.properties or {}).items():
+                    if not isinstance(property_, dm.MappedPropertyApply):
+                        continue
+                    container_entity = ContainerEntity.from_id(property_.container)
+                    view_properties_by_container_properties[
+                        (container_entity, property_.container_property_identifier)
+                    ].append((view_entity, property_id))
+
+        return view_properties_by_container_properties
+
+    @staticmethod
+    def _get_parent_view_property_by_container_property(
+        path_by_view, view_by_container_properties: dict[tuple[ContainerEntity, str], list[tuple[ViewEntity, str]]]
+    ) -> dict[tuple[ContainerEntity, str], str]:
+        parent_name_by_container_property: dict[tuple[ContainerEntity, str], str] = {}
+        for (container, container_property), view_properties in view_by_container_properties.items():
+            if len(view_properties) == 1:
+                continue
+            # Shortest path is the parent
+            _, prop_name = min(view_properties, key=lambda prop: len(path_by_view[prop[0]]))
+            parent_name_by_container_property[(container, container_property)] = prop_name
+        return parent_name_by_container_property

cognite/neat/_rules/transformers/_pipelines.py

@@ -23,7 +23,7 @@ class ImporterPipeline(RulesPipeline[InputRules, VerifiedRules]):
 
     @classmethod
     def _create_pipeline(cls, importer: BaseImporter[InputRules], role: RoleTypes | None = None) -> "ImporterPipeline":
-        items: list[RulesTransformer] = [VerifyAnyRules(errors="continue")]
+        items: list[RulesTransformer] = [VerifyAnyRules(errors="continue", validate=True)]
         if role is not None:
             out_cls = VERIFIED_RULES_BY_ROLE[role]
             items.append(ConvertToRules(out_cls))

cognite/neat/_rules/transformers/_verification.py

@@ -1,7 +1,7 @@
 from abc import ABC
 from typing import Any, Literal
 
-from cognite.neat._issues import IssueList, NeatError, NeatWarning, catch_issues
+from cognite.neat._issues import IssueList, MultiValueError, NeatError, NeatWarning, catch_issues
 from cognite.neat._issues.errors import NeatTypeError
 from cognite.neat._rules._shared import (
     InputRules,
@@ -18,6 +18,8 @@ from cognite.neat._rules.models import (
     InformationInputRules,
     InformationRules,
 )
+from cognite.neat._rules.models.dms import DMSValidation
+from cognite.neat._rules.models.information import InformationValidation
 
 from ._base import RulesTransformer
 
@@ -26,10 +28,11 @@ class VerificationTransformer(RulesTransformer[T_InputRules, T_VerifiedRules], A
     """Base class for all verification transformers."""
 
     _rules_cls: type[T_VerifiedRules]
+    _validation_cls: type
 
-    def __init__(self, errors: Literal["raise", "continue"], post_validate: bool = True) -> None:
+    def __init__(self, errors: Literal["raise", "continue"], validate: bool = True) -> None:
         self.errors = errors
-        self.post_validate = post_validate
+        self.validate = validate
 
     def transform(self, rules: T_InputRules | OutRules[T_InputRules]) -> MaybeRules[T_VerifiedRules]:
         issues = IssueList()
@@ -41,8 +44,17 @@ class VerificationTransformer(RulesTransformer[T_InputRules, T_VerifiedRules], A
         with catch_issues(issues, NeatError, NeatWarning, error_args) as future:
             rules_cls = self._get_rules_cls(in_)
             dumped = in_.dump()
-            dumped["post_validate"] = self.post_validate
             verified_rules = rules_cls.model_validate(dumped)  # type: ignore[assignment]
+            if self.validate:
+                validation_cls = self._get_validation_cls(verified_rules)  # type: ignore[arg-type]
+                validation_issues = validation_cls(verified_rules).validate()
+                # We need to trigger warnings are raise exceptions such that they are caught by the context manager
+                # and processed with the read context
+                if validation_issues.warnings:
+                    validation_issues.trigger_warnings()
+                if validation_issues.has_errors:
+                    verified_rules = None
+                    raise MultiValueError(validation_issues.errors)
 
         if (future.result == "failure" or issues.has_errors or verified_rules is None) and self.errors == "raise":
             raise issues.as_errors()
@@ -54,17 +66,22 @@ class VerificationTransformer(RulesTransformer[T_InputRules, T_VerifiedRules], A
     def _get_rules_cls(self, in_: T_InputRules) -> type[T_VerifiedRules]:
         return self._rules_cls
 
+    def _get_validation_cls(self, rules: T_VerifiedRules) -> type:
+        return self._validation_cls
+
 
 class VerifyDMSRules(VerificationTransformer[DMSInputRules, DMSRules]):
     """Class to verify DMS rules."""
 
     _rules_cls = DMSRules
+    _validation_cls = DMSValidation
 
 
 class VerifyInformationRules(VerificationTransformer[InformationInputRules, InformationRules]):
     """Class to verify Information rules."""
 
     _rules_cls = InformationRules
+    _validation_cls = InformationValidation
 
 
 class VerifyAnyRules(VerificationTransformer[InputRules, VerifiedRules]):
@@ -77,3 +94,11 @@ class VerifyAnyRules(VerificationTransformer[InputRules, VerifiedRules]):
             return DMSRules
         else:
             raise NeatTypeError(f"Unsupported rules type: {type(in_)}")
+
+    def _get_validation_cls(self, rules: T_VerifiedRules) -> type:
+        if isinstance(rules, InformationRules):
+            return InformationValidation
+        elif isinstance(rules, DMSRules):
+            return DMSValidation
+        else:
+            raise NeatTypeError(f"Unsupported rules type: {type(rules)}")