[NEAT-463] Core Extension (#672)

* docs; started on new tutorial

* docs; finished shell of tutorial

* refactor: started on prepare methods

* refactor: fix extension method

* fix: handle property ID

* refactor: implementd to excel

* fix: maximum column width

* refactor; remove cognite prefix

* refactor: added new prefix

* refactor: added view reduction

* fix: drop 3D property

* refactor; neat session use in memory by default

* refactor: setup dms analysis

* refactor: expose new analysis

* refactor: fix include all properties

* docs: finished doc tutorial

* docs; clear hidden output

* refactor: review feedback

* Cleanup dependencies (#674)

* build: moved around dependencies

* style: happy mypy

* build: pin python-multipart

* docs; document

cognite/neat/_constants.py

@@ -1,11 +1,16 @@
 from pathlib import Path
 
+from cognite.client.data_classes.data_modeling.ids import DataModelId
 from rdflib import DCTERMS, OWL, RDF, RDFS, SKOS, XSD, Namespace, URIRef
 
 from cognite import neat
 
 PACKAGE_DIRECTORY = Path(neat.__file__).parent
-
+COGNITE_MODELS = (
+    DataModelId("cdf_cdm", "CogniteCore", "v1"),
+    DataModelId("cdf_idm", "CogniteProcessIndustries", "v1"),
+)
+DMS_LISTABLE_PROPERTY_LIMIT = 1000
 
 EXAMPLE_RULES = PACKAGE_DIRECTORY / "_rules" / "examples"
 EXAMPLE_GRAPHS = PACKAGE_DIRECTORY / "_graph" / "examples"

cognite/neat/_issues/_base.py

@@ -11,7 +11,7 @@
 from warnings import WarningMessage
 
 import pandas as pd
-from cognite.client.data_classes.data_modeling import ContainerId, ViewId
+from cognite.client.data_classes.data_modeling import ContainerId, PropertyId, ViewId
 from pydantic_core import ErrorDetails
 
 from cognite.neat._utils.spreadsheet import SpreadsheetRead
@@ -121,6 +121,8 @@ def _dump_value(cls, value: Any) -> list | int | bool | float | str | dict:
             return value.dump(camel_case=True, include_type=True)
         elif isinstance(value, Entity):
             return value.dump()
+        elif isinstance(value, PropertyId):
+            return value.dump(camel_case=True)
         raise ValueError(f"Unsupported type: {type(value)}")
 
     @classmethod
@@ -151,7 +153,7 @@ def _load_values(cls, neat_issue_cls: "type[NeatIssue]", data: dict[str, Any]) -
         return neat_issue_cls(**args)
 
     @classmethod
-    def _load_value(cls, type_: type, value: Any) -> Any:
+    def _load_value(cls, type_: Any, value: Any) -> Any:
         from cognite.neat._rules.models.entities import Entity
 
         if isinstance(type_, UnionType) or get_origin(type_) is UnionType:
@@ -167,6 +169,8 @@ def _load_value(cls, type_: type, value: Any) -> Any:
             return tuple(cls._load_value(subtype, item) for item in value)
         elif type_ is ViewId:
             return ViewId.load(value)
+        elif type_ is PropertyId:
+            return PropertyId.load(value)
         elif type_ is ContainerId:
             return ContainerId.load(value)
         elif inspect.isclass(type_) and issubclass(type_, Entity):

cognite/neat/_rules/analysis/__init__.py

@@ -1,6 +1,5 @@
 from ._asset import AssetAnalysis
-from ._information import (
-    InformationAnalysis,
-)
+from ._dms import DMSAnalysis
+from ._information import InformationAnalysis
 
-__all__ = ["InformationAnalysis", "AssetAnalysis"]
+__all__ = ["InformationAnalysis", "AssetAnalysis", "DMSAnalysis"]

cognite/neat/_rules/analysis/_dms.py

@@ -0,0 +1,43 @@
+from cognite.neat._constants import DMS_LISTABLE_PROPERTY_LIMIT
+from cognite.neat._rules.models.dms import DMSProperty, DMSRules, DMSView
+from cognite.neat._rules.models.entities import ReferenceEntity, ViewEntity
+
+from ._base import BaseAnalysis
+
+
+class DMSAnalysis(BaseAnalysis[DMSRules, DMSView, DMSProperty, ViewEntity, str]):
+    """Assumes analysis over only the complete schema"""
+
+    def _get_classes(self) -> list[DMSView]:
+        return list(self.rules.views)
+
+    def _get_properties(self) -> list[DMSProperty]:
+        return list(self.rules.properties)
+
+    def _get_reference(self, class_or_property: DMSView | DMSProperty) -> ReferenceEntity | None:
+        return class_or_property.reference if isinstance(class_or_property.reference, ReferenceEntity) else None
+
+    def _get_cls_entity(self, class_: DMSView | DMSProperty) -> ViewEntity:
+        return class_.view
+
+    def _get_prop_entity(self, property_: DMSProperty) -> str:
+        return property_.property_
+
+    def _get_cls_parents(self, class_: DMSView) -> list[ViewEntity] | None:
+        return list(class_.implements) if class_.implements else None
+
+    def _get_reference_rules(self) -> DMSRules | None:
+        return self.rules.reference
+
+    @classmethod
+    def _set_cls_entity(cls, property_: DMSProperty, class_: ViewEntity) -> None:
+        property_.view = class_
+
+    def _get_object(self, property_: DMSProperty) -> ViewEntity | None:
+        return property_.value_type if isinstance(property_.value_type, ViewEntity) else None
+
+    def _get_max_occurrence(self, property_: DMSProperty) -> int | float | None:
+        return DMS_LISTABLE_PROPERTY_LIMIT if property_.is_list else 1
+
+    def subset_rules(self, desired_classes: set[ViewEntity]) -> DMSRules:
+        raise NotImplementedError()

cognite/neat/_rules/exporters/_rules2excel.py

@@ -27,6 +27,8 @@
 
 from ._base import BaseExporter
 
+MAX_COLUMN_WIDTH = 70.0
+
 
 class ExcelExporter(BaseExporter[VerifiedRules, Workbook]):
     """Export rules to Excel.
@@ -251,7 +253,9 @@ def _adjust_column_widths(cls, workbook: Workbook) -> None:
                     selected_column = column_cells[1]
 
                 current = sheet.column_dimensions[selected_column.column_letter].width or (max_length + 0.5)
-                sheet.column_dimensions[selected_column.column_letter].width = max(current, max_length + 0.5)
+                sheet.column_dimensions[selected_column.column_letter].width = min(
+                    max(current, max_length + 0.5), MAX_COLUMN_WIDTH
+                )
         return None
 
 

cognite/neat/_rules/models/_base_input.py

@@ -64,14 +64,15 @@ def _type_by_field_name(cls) -> dict[str, type]:
             if isinstance(type_, str) and type_.startswith(cls.__name__):
                 type_ = cls
 
+            candidate: type
             if is_dataclass(type_):
-                candidate = type_
+                candidate = type_  # type: ignore[assignment]
             elif isinstance(type_, GenericAlias) and type_.__origin__ is list and is_dataclass(type_.__args__[0]):
-                candidate = type_.__args__[0]
+                candidate = type_.__args__[0]  # type: ignore[assignment]
 
             # this handles prefixes
             elif isinstance(type_, GenericAlias) and type_.__origin__ is dict:
-                candidate = type_
+                candidate = type_  # type: ignore[assignment]
             else:
                 continue
 

cognite/neat/_rules/models/dms/_schema.py

@@ -461,7 +461,7 @@ def load(cls, data: str | dict[str, list[Any]], context: dict[str, list[Path]] |
                         )
                 else:
                     try:
-                        loaded[attr.name] = attr.type.load(items)
+                        loaded[attr.name] = attr.type.load(items)  # type: ignore[union-attr]
                     except Exception as e:
                         loaded[attr.name] = cls._load_individual_resources(
                             items, attr, str(e), context.get(attr.name, [])
@@ -470,14 +470,15 @@ def load(cls, data: str | dict[str, list[Any]], context: dict[str, list[Path]] |
 
     @classmethod
     def _load_individual_resources(cls, items: list, attr: Field, trigger_error: str, resource_context) -> list[Any]:
-        resources = attr.type([])
-        if not hasattr(attr.type, "_RESOURCE"):
+        type_ = cast(type, attr.type)
+        resources = type_([])
+        if not hasattr(type_, "_RESOURCE"):
             warnings.warn(
-                FileTypeUnexpectedWarning(Path("UNKNOWN"), frozenset([attr.type.__name__]), trigger_error), stacklevel=2
+                FileTypeUnexpectedWarning(Path("UNKNOWN"), frozenset([type_.__name__]), trigger_error), stacklevel=2
             )
             return resources
         # Fallback to load individual resources.
-        single_cls = attr.type._RESOURCE
+        single_cls = type_._RESOURCE
         for no, item in enumerate(items):
             try:
                 loaded_instance = single_cls.load(item)

cognite/neat/_rules/transformers/__init__.py

@@ -5,7 +5,9 @@
     DMSToInformation,
     InformationToAsset,
     InformationToDMS,
+    ReduceCogniteModel,
     ToCompliantEntities,
+    ToExtension,
 )
 from ._mapping import MapOneToOne, RuleMapper
 from ._pipelines import ImporterPipeline
@@ -27,4 +29,6 @@
     "MapOneToOne",
     "ToCompliantEntities",
     "RuleMapper",
+    "ToExtension",
+    "ReduceCogniteModel",
 ]

cognite/neat/_rules/transformers/_converters.py

@@ -2,19 +2,23 @@
 import warnings
 from abc import ABC, abstractmethod
 from collections import Counter, defaultdict
-from collections.abc import Collection
+from collections.abc import Collection, Mapping
 from datetime import date, datetime
 from typing import Literal, TypeVar, cast
 
 from cognite.client.data_classes import data_modeling as dms
+from cognite.client.data_classes.data_modeling import DataModelId, DataModelIdentifier, ViewId
 from rdflib import Namespace
 
-from cognite.neat._constants import DMS_CONTAINER_PROPERTY_SIZE_LIMIT
+from cognite.neat._constants import COGNITE_MODELS, DMS_CONTAINER_PROPERTY_SIZE_LIMIT
+from cognite.neat._issues.errors import NeatValueError
 from cognite.neat._issues.warnings.user_modeling import ParentInDifferentSpaceWarning
 from cognite.neat._rules._constants import EntityTypes
 from cognite.neat._rules._shared import InputRules, JustRules, OutRules, VerifiedRules
+from cognite.neat._rules.analysis import DMSAnalysis
 from cognite.neat._rules.models import (
     AssetRules,
+    DMSInputRules,
     DMSRules,
     DomainRules,
     ExtensionCategory,
@@ -220,6 +224,146 @@ def _transform(self, rules: VerifiedRules) -> VerifiedRules:
         raise ValueError(f"Unsupported conversion from {type(rules)} to {self._out_cls}")
 
 
+_T_Entity = TypeVar("_T_Entity", bound=ClassEntity | ViewEntity)
+
+
+class ToExtension(RulesTransformer[DMSRules, DMSRules]):
+    def __init__(
+        self,
+        new_model_id: DataModelIdentifier,
+        org_name: str | None = None,
+        mode: Literal["composition"] = "composition",
+    ):
+        self.new_model_id = DataModelId.load(new_model_id)
+        self.org_name = org_name
+        self.mode = mode
+
+    def transform(self, rules: DMSRules | OutRules[DMSRules]) -> JustRules[DMSRules]:
+        # Copy to ensure immutability
+        verified = self._to_rules(rules)
+        if self.org_name is None and verified.metadata.as_data_model_id() in COGNITE_MODELS:
+            raise NeatValueError(f"Prefix is required when extending {verified.metadata.as_data_model_id()}")
+        source_id = verified.metadata.as_data_model_id()
+
+        dump = verified.dump()
+        dump["metadata"]["schema_"] = SchemaCompleteness.partial.value
+        dump["metadata"]["space"] = self.new_model_id.space
+        dump["metadata"]["external_id"] = self.new_model_id.external_id
+        if self.new_model_id.version is not None:
+            dump["metadata"]["version"] = self.new_model_id.version
+        # Serialize and deserialize to set the new space and external_id
+        # as the default values for the new model.
+        new_model = DMSRules.model_validate(DMSInputRules.load(dump).dump())
+
+        # Write back the original space and external_id for the container of the new model.
+        for prop in new_model.properties:
+            if prop.container and prop.container.space == self.new_model_id.space:
+                prop.container = ContainerEntity(
+                    space=source_id.space,
+                    externalId=prop.container.suffix,
+                )
+
+        if self.mode == "composition":
+            new_model.containers = None
+            for view in new_model.views:
+                view.implements = None
+
+        if source_id in COGNITE_MODELS:
+            # Remove CognitePrefixes.
+            for prop in new_model.properties:
+                prop.view = self._remove_cognite_prefix(prop.view)
+                prop.class_ = self._remove_cognite_prefix(prop.class_)
+                if isinstance(prop.value_type, ViewEntity):
+                    prop.value_type = self._remove_cognite_prefix(prop.value_type)
+            for view in new_model.views:
+                view.view = self._remove_cognite_prefix(view.view)
+                view.class_ = self._remove_cognite_prefix(view.class_)
+
+        return JustRules(new_model)
+
+    def _remove_cognite_prefix(self, entity: _T_Entity) -> _T_Entity:
+        new_suffix = entity.suffix.replace("Cognite", self.org_name or "")
+        if isinstance(entity, ViewEntity):
+            return ViewEntity(space=entity.space, externalId=new_suffix, version=entity.version)  # type: ignore[return-value]
+        elif isinstance(entity, ClassEntity):
+            return ClassEntity(prefix=entity.prefix, suffix=new_suffix, version=entity.version)  # type: ignore[return-value]
+        raise ValueError(f"Unsupported entity type: {type(entity)}")
+
+
+class ReduceCogniteModel(RulesTransformer[DMSRules, DMSRules]):
+    _ASSET_VIEW = ViewId("cdf_cdm", "CogniteAsset", "v1")
+    _VIEW_BY_COLLECTION: Mapping[Literal["3D", "Annotation", "BaseViews"], frozenset[ViewId]] = {
+        "3D": frozenset(
+            {
+                ViewId("cdf_cdm", "Cognite3DModel", "v1"),
+                ViewId("cdf_cdm", "Cognite3DObject", "v1"),
+                ViewId("cdf_cdm", "Cognite3DRevision", "v1"),
+                ViewId("cdf_cdm", "Cognite3DTransformation", "v1"),
+                ViewId("cdf_cdm", "Cognite360Image", "v1"),
+                ViewId("cdf_cdm", "Cognite360ImageAnnotation", "v1"),
+                ViewId("cdf_cdm", "Cognite360ImageCollection", "v1"),
+                ViewId("cdf_cdm", "Cognite360ImageModel", "v1"),
+                ViewId("cdf_cdm", "Cognite360ImageStation", "v1"),
+                ViewId("cdf_cdm", "CogniteCADModel", "v1"),
+                ViewId("cdf_cdm", "CogniteCADNode", "v1"),
+                ViewId("cdf_cdm", "CogniteCADRevision", "v1"),
+                ViewId("cdf_cdm", "CogniteCubeMap", "v1"),
+                ViewId("cdf_cdm", "CognitePointCloudModel", "v1"),
+                ViewId("cdf_cdm", "CognitePointCloudRevision", "v1"),
+                ViewId("cdf_cdm", "CognitePointCloudVolume", "v1"),
+            }
+        ),
+        "Annotation": frozenset(
+            {
+                ViewId("cdf_cdm", "CogniteAnnotation", "v1"),
+                ViewId("cdf_cdm", "CogniteDiagramAnnotation", "v1"),
+            }
+        ),
+        "BaseViews": frozenset(
+            {
+                ViewId("cdf_cdm", "CogniteDescribable", "v1"),
+                ViewId("cdf_cdm", "CogniteSchedulable", "v1"),
+                ViewId("cdf_cdm", "CogniteSourceable", "v1"),
+                ViewId("cdf_cdm", "CogniteVisualizable", "v1"),
+            }
+        ),
+    }
+
+    def __init__(self, drop: Collection[Literal["3D", "Annotation", "BaseViews"]]):
+        self.drop = drop
+
+    def transform(self, rules: DMSRules | OutRules[DMSRules]) -> JustRules[DMSRules]:
+        verified = self._to_rules(rules)
+        if verified.metadata.as_data_model_id() not in COGNITE_MODELS:
+            raise NeatValueError(f"Can only reduce Cognite Data Models, not {verified.metadata.as_data_model_id()}")
+        if invalid := (set(self.drop) - set(self._VIEW_BY_COLLECTION.keys())):
+            raise NeatValueError(f"Invalid drop values: {invalid}. Expected {set(self._VIEW_BY_COLLECTION)}")
+
+        exclude_views = {view for collection in self.drop for view in self._VIEW_BY_COLLECTION[collection]}
+        new_model = verified.model_copy(deep=True)
+
+        properties_by_view = DMSAnalysis(new_model).classes_with_properties(consider_inheritance=True)
+
+        new_model.views = SheetList[DMSView](
+            [view for view in new_model.views if view.view.as_id() not in exclude_views]
+        )
+        new_properties = SheetList[DMSProperty]()
+        for view in new_model.views:
+            for prop in properties_by_view[view.view]:
+                if self._is_asset_3D_property(prop):
+                    # We filter out the 3D property of asset
+                    continue
+                new_properties.append(prop)
+        new_model.properties = new_properties
+
+        return JustRules(new_model)
+
+    def _is_asset_3D_property(self, prop: DMSProperty) -> bool:
+        if "3D" not in self.drop:
+            return False
+        return prop.view.as_id() == self._ASSET_VIEW and prop.property_ == "object3D"
+
+
 class _InformationRulesConverter:
     def __init__(self, information: InformationRules):
         self.rules = information

cognite/neat/_session/_base.py

@@ -25,7 +25,7 @@ class NeatSession:
     def __init__(
         self,
         client: CogniteClient | None = None,
-        storage: Literal["memory", "oxigraph"] = "oxigraph",
+        storage: Literal["memory", "oxigraph"] = "memory",
         verbose: bool = True,
     ) -> None:
         self._client = client