cognite-neat 0.98.0__py3-none-any.whl → 0.99.1__py3-none-any.whl

cognite/neat/_client/data_classes/schema.py

@@ -0,0 +1,495 @@
+import sys
+import warnings
+import zipfile
+from collections import ChainMap
+from collections.abc import Iterable, MutableMapping
+from dataclasses import Field, dataclass, field, fields
+from pathlib import Path
+from typing import Any, ClassVar, Literal, cast
+
+import yaml
+from cognite.client import data_modeling as dm
+from cognite.client.data_classes import DatabaseWrite, TransformationWrite
+from cognite.client.data_classes.data_modeling.views import (
+    ReverseDirectRelationApply,
+    SingleEdgeConnection,
+    SingleEdgeConnectionApply,
+    SingleReverseDirectRelation,
+    SingleReverseDirectRelationApply,
+    ViewProperty,
+    ViewPropertyApply,
+)
+
+from cognite.neat._client.data_classes.data_modeling import (
+    CogniteResourceDict,
+    ContainerApplyDict,
+    NodeApplyDict,
+    SpaceApplyDict,
+    ViewApplyDict,
+)
+from cognite.neat._issues.errors import (
+    NeatYamlError,
+)
+from cognite.neat._issues.warnings import (
+    FileTypeUnexpectedWarning,
+    ResourcesDuplicatedWarning,
+)
+from cognite.neat._utils.text import to_camel
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+
+@dataclass
+class DMSSchema:
+    data_model: dm.DataModelApply | None = None
+    spaces: SpaceApplyDict = field(default_factory=SpaceApplyDict)
+    views: ViewApplyDict = field(default_factory=ViewApplyDict)
+    containers: ContainerApplyDict = field(default_factory=ContainerApplyDict)
+    node_types: NodeApplyDict = field(default_factory=NodeApplyDict)
+
+    _FIELD_NAME_BY_RESOURCE_TYPE: ClassVar[dict[str, str]] = {
+        "container": "containers",
+        "view": "views",
+        "datamodel": "data_model",
+        "space": "spaces",
+        "node": "node_types",
+    }
+
+    @classmethod
+    def from_directory(cls, directory: str | Path) -> Self:
+        """Load a schema from a directory containing YAML files.
+
+        The directory is expected to follow the Cognite-Toolkit convention
+        where each file is named as `resource_type.resource_name.yaml`.
+        """
+        data, context = cls._read_directory(Path(directory))
+        return cls.load(data, context)
+
+    @classmethod
+    def _read_directory(cls, directory: Path) -> tuple[dict[str, list[Any]], dict[str, list[Path]]]:
+        data: dict[str, Any] = {}
+        context: dict[str, list[Path]] = {}
+        for yaml_file in directory.rglob("*.yaml"):
+            if "." in yaml_file.stem:
+                resource_type = yaml_file.stem.rsplit(".", 1)[-1]
+                if attr_name := cls._FIELD_NAME_BY_RESOURCE_TYPE.get(resource_type):
+                    data.setdefault(attr_name, [])
+                    context.setdefault(attr_name, [])
+                    try:
+                        loaded = yaml.safe_load(yaml_file.read_text())
+                    except Exception as e:
+                        warnings.warn(
+                            FileTypeUnexpectedWarning(yaml_file, frozenset([".yaml", ".yml"]), str(e)), stacklevel=2
+                        )
+                        continue
+
+                    if isinstance(loaded, list):
+                        data[attr_name].extend(loaded)
+                        context[attr_name].extend([yaml_file] * len(loaded))
+                    else:
+                        data[attr_name].append(loaded)
+                        context[attr_name].append(yaml_file)
+        return data, context
+
+    def to_directory(
+        self,
+        directory: str | Path,
+        exclude: set[str] | None = None,
+        new_line: str | None = "\n",
+        encoding: str | None = "utf-8",
+    ) -> None:
+        """Save the schema to a directory as YAML files. This is compatible with the Cognite-Toolkit convention.
+
+        Args:
+            directory (str | Path): The directory to save the schema to.
+            exclude (set[str]): A set of attributes to exclude from the output.
+            new_line (str): The line endings to use in the output files. Defaults to "\n".
+            encoding (str): The encoding to use in the output files. Defaults to "utf-8".
+        """
+        path_dir = Path(directory)
+        exclude_set = exclude or set()
+        data_models = path_dir / "data_models"
+        data_models.mkdir(parents=True, exist_ok=True)
+        if "spaces" not in exclude_set:
+            for space in self.spaces.values():
+                (data_models / f"{space.space}.space.yaml").write_text(
+                    space.dump_yaml(), newline=new_line, encoding=encoding
+                )
+        if "data_models" not in exclude_set and self.data_model:
+            (data_models / f"{self.data_model.external_id}.datamodel.yaml").write_text(
+                self.data_model.dump_yaml(), newline=new_line, encoding=encoding
+            )
+        if "views" not in exclude_set and self.views:
+            view_dir = data_models / "views"
+            view_dir.mkdir(parents=True, exist_ok=True)
+            for view in self.views.values():
+                (view_dir / f"{view.external_id}.view.yaml").write_text(
+                    view.dump_yaml(), newline=new_line, encoding=encoding
+                )
+        if "containers" not in exclude_set and self.containers:
+            container_dir = data_models / "containers"
+            container_dir.mkdir(parents=True, exist_ok=True)
+            for container in self.containers.values():
+                (container_dir / f"{container.external_id}.container.yaml").write_text(
+                    container.dump_yaml(), newline=new_line, encoding=encoding
+                )
+        if "node_types" not in exclude_set and self.node_types:
+            node_dir = data_models / "nodes"
+            node_dir.mkdir(parents=True, exist_ok=True)
+            for node in self.node_types.values():
+                (node_dir / f"{node.external_id}.node.yaml").write_text(
+                    node.dump_yaml(), newline=new_line, encoding=encoding
+                )
+
+    @classmethod
+    def from_zip(cls, zip_file: str | Path) -> Self:
+        """Load a schema from a ZIP file containing YAML files.
+
+        The ZIP file is expected to follow the Cognite-Toolkit convention
+        where each file is named as `resource_type.resource_name.yaml`.
+        """
+        data, context = cls._read_zip(Path(zip_file))
+        return cls.load(data, context)
+
+    @classmethod
+    def _read_zip(cls, zip_file: Path) -> tuple[dict[str, list[Any]], dict[str, list[Path]]]:
+        data: dict[str, list[Any]] = {}
+        context: dict[str, list[Path]] = {}
+        with zipfile.ZipFile(zip_file, "r") as zip_ref:
+            for file_info in zip_ref.infolist():
+                if file_info.filename.endswith(".yaml"):
+                    if "/" not in file_info.filename:
+                        continue
+                    filename = Path(file_info.filename.split("/")[-1])
+                    if "." not in filename.stem:
+                        continue
+                    resource_type = filename.stem.rsplit(".", 1)[-1]
+                    if attr_name := cls._FIELD_NAME_BY_RESOURCE_TYPE.get(resource_type):
+                        data.setdefault(attr_name, [])
+                        context.setdefault(attr_name, [])
+                        try:
+                            loaded = yaml.safe_load(zip_ref.read(file_info).decode())
+                        except Exception as e:
+                            warnings.warn(
+                                FileTypeUnexpectedWarning(filename, frozenset([".yaml", ".yml"]), str(e)), stacklevel=2
+                            )
+                            continue
+                        if isinstance(loaded, list):
+                            data[attr_name].extend(loaded)
+                            context[attr_name].extend([filename] * len(loaded))
+                        else:
+                            data[attr_name].append(loaded)
+                            context[attr_name].append(filename)
+        return data, context
+
+    def to_zip(self, zip_file: str | Path, exclude: set[str] | None = None) -> None:
+        """Save the schema to a ZIP file as YAML files. This is compatible with the Cognite-Toolkit convention.
+
+        Args:
+            zip_file (str | Path): The ZIP file to save the schema to.
+            exclude (set[str]): A set of attributes to exclude from the output.
+        """
+        exclude_set = exclude or set()
+        with zipfile.ZipFile(zip_file, "w") as zip_ref:
+            if "spaces" not in exclude_set:
+                for space in self.spaces.values():
+                    zip_ref.writestr(f"data_models/{space.space}.space.yaml", space.dump_yaml())
+            if "data_models" not in exclude_set and self.data_model:
+                zip_ref.writestr(
+                    f"data_models/{self.data_model.external_id}.datamodel.yaml", self.data_model.dump_yaml()
+                )
+            if "views" not in exclude_set:
+                for view in self.views.values():
+                    zip_ref.writestr(f"data_models/views/{view.external_id}.view.yaml", view.dump_yaml())
+            if "containers" not in exclude_set:
+                for container in self.containers.values():
+                    zip_ref.writestr(
+                        f"data_models/containers{container.external_id}.container.yaml", container.dump_yaml()
+                    )
+            if "node_types" not in exclude_set:
+                for node in self.node_types.values():
+                    zip_ref.writestr(f"data_models/nodes/{node.external_id}.node.yaml", node.dump_yaml())
+
+    @classmethod
+    def load(cls, data: str | dict[str, list[Any]], context: dict[str, list[Path]] | None = None) -> Self:
+        """Loads a schema from a dictionary or a YAML or JSON formatted string.
+
+        Args:
+            data: The data to load the schema from. This can be a dictionary, a YAML or JSON formatted string.
+            context: This provides linage for where the data was loaded from. This is used in Warnings
+                if a single item fails to load.
+
+        Returns:
+            DMSSchema: The loaded schema.
+        """
+        context = context or {}
+        if isinstance(data, str):
+            # YAML is a superset of JSON, so we can use the same parser
+            try:
+                data_dict = yaml.safe_load(data)
+            except Exception as e:
+                raise NeatYamlError(str(e)) from None
+            if not isinstance(data_dict, dict) and all(isinstance(v, list) for v in data_dict.values()):
+                raise NeatYamlError(f"Invalid data structure: {type(data)}", "dict[str, list[Any]]") from None
+        else:
+            data_dict = data
+        loaded: dict[str, Any] = {}
+        for attr in fields(cls):
+            if items := data_dict.get(attr.name) or data_dict.get(to_camel(attr.name)):
+                if attr.name == "data_model":
+                    if isinstance(items, list) and len(items) > 1:
+                        try:
+                            data_model_ids = [dm.DataModelId.load(item) for item in items]
+                        except Exception as e:
+                            data_model_file = context.get(attr.name, [Path("UNKNOWN")])[0]
+                            warnings.warn(
+                                FileTypeUnexpectedWarning(
+                                    data_model_file, frozenset([dm.DataModelApply.__name__]), str(e)
+                                ),
+                                stacklevel=2,
+                            )
+                        else:
+                            warnings.warn(
+                                ResourcesDuplicatedWarning(
+                                    frozenset(data_model_ids),
+                                    "data model",
+                                    "Will use the first DataModel.",
+                                ),
+                                stacklevel=2,
+                            )
+                    item = items[0] if isinstance(items, list) else items
+                    try:
+                        loaded[attr.name] = dm.DataModelApply.load(item)
+                    except Exception as e:
+                        data_model_file = context.get(attr.name, [Path("UNKNOWN")])[0]
+                        warnings.warn(
+                            FileTypeUnexpectedWarning(data_model_file, frozenset([dm.DataModelApply.__name__]), str(e)),
+                            stacklevel=2,
+                        )
+                else:
+                    try:
+                        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, [])
+                        )
+        return cls(**loaded)
+
+    @classmethod
+    def _load_individual_resources(cls, items: list, attr: Field, trigger_error: str, resource_context) -> list[Any]:
+        type_ = cast(type, attr.type)
+        resources = type_([])
+        if not hasattr(type_, "_RESOURCE"):
+            warnings.warn(
+                FileTypeUnexpectedWarning(Path("UNKNOWN"), frozenset([type_.__name__]), trigger_error), stacklevel=2
+            )
+            return resources
+        # Fallback to load individual resources.
+        single_cls = type_._RESOURCE
+        for no, item in enumerate(items):
+            try:
+                loaded_instance = single_cls.load(item)
+            except Exception as e:
+                try:
+                    filepath = resource_context[no]
+                except IndexError:
+                    filepath = Path("UNKNOWN")
+                # We use repr(e) instead of str(e) to include the exception type in the warning message
+                warnings.warn(
+                    FileTypeUnexpectedWarning(filepath, frozenset([single_cls.__name__]), repr(e)), stacklevel=2
+                )
+            else:
+                resources.append(loaded_instance)
+        return resources
+
+    def dump(self, camel_case: bool = True, sort: bool = True) -> dict[str, Any]:
+        """Dump the schema to a dictionary that can be serialized to JSON.
+
+        Args:
+            camel_case (bool): If True, the keys in the output dictionary will be in camel case.
+            sort (bool): If True, the items in the output dictionary will be sorted by their ID.
+                This is useful for deterministic output which is useful for comparing schemas.
+
+        Returns:
+            dict: The schema as a dictionary.
+        """
+        output: dict[str, Any] = {}
+        cls_fields = sorted(fields(self), key=lambda f: f.name) if sort else fields(self)
+        for attr in cls_fields:
+            if items := getattr(self, attr.name):
+                key = to_camel(attr.name) if camel_case else attr.name
+                if isinstance(items, CogniteResourceDict):
+                    if sort:
+                        output[key] = [
+                            item.dump(camel_case) for item in sorted(items.values(), key=self._to_sortable_identifier)
+                        ]
+                    else:
+                        output[key] = items.dump(camel_case)
+                else:
+                    output[key] = items.dump(camel_case=camel_case)
+        return output
+
+    @classmethod
+    def _to_sortable_identifier(cls, item: Any) -> str | tuple[str, str] | tuple[str, str, str]:
+        if isinstance(item, dm.ContainerApply | dm.ViewApply | dm.DataModelApply | dm.NodeApply):
+            identifier = item.as_id().as_tuple()
+            if len(identifier) == 3 and identifier[2] is None:
+                return identifier[:2]  # type: ignore[misc]
+            return cast(tuple[str, str] | tuple[str, str, str], identifier)
+        elif isinstance(item, dm.SpaceApply):
+            return item.space
+        elif isinstance(item, TransformationWrite):
+            return item.external_id or ""
+        elif isinstance(item, DatabaseWrite):
+            return item.name or ""
+        else:
+            raise ValueError(f"Cannot sort item of type {type(item)}")
+
+    def referenced_spaces(self, include_indirect_references: bool = True) -> set[str]:
+        """Get the spaces referenced by the schema.
+
+        Args:
+            include_indirect_references (bool): If True, the spaces referenced by as view.implements, and
+                view.referenced_containers will be included in the output.
+        Returns:
+            set[str]: The spaces referenced by the schema.
+        """
+        referenced_spaces = {view.space for view in self.views.values()}
+        referenced_spaces |= {container.space for container in self.containers.values()}
+        if include_indirect_references:
+            referenced_spaces |= {
+                container.space for view in self.views.values() for container in view.referenced_containers()
+            }
+            referenced_spaces |= {parent.space for view in self.views.values() for parent in view.implements or []}
+        referenced_spaces |= {node.space for node in self.node_types.values()}
+        if self.data_model:
+            referenced_spaces |= {self.data_model.space}
+            referenced_spaces |= {view.space for view in self.data_model.views or []}
+        referenced_spaces |= {s.space for s in self.spaces.values()}
+        return referenced_spaces
+
+    def referenced_container(self) -> set[dm.ContainerId]:
+        referenced_containers = {
+            container for view in self.views.values() for container in view.referenced_containers()
+        }
+        referenced_containers |= set(self.containers.keys())
+        return referenced_containers
+
+    def as_read_model(self) -> dm.DataModel[dm.View]:
+        if self.data_model is None:
+            raise ValueError("Data model is not defined")
+        all_containers = self.containers.copy()
+        all_views = self.views.copy()
+        views: list[dm.View] = []
+        for view in self.views.values():
+            referenced_containers = ContainerApplyDict()
+            properties: dict[str, ViewProperty] = {}
+            # ChainMap is used to merge properties from the view and its parents
+            # Note that the order of the ChainMap is important, as the first dictionary has the highest priority
+            # So if a child and parent have the same property, the child property will be used.
+            write_properties = ChainMap(view.properties, *(all_views[v].properties for v in view.implements or []))  # type: ignore[arg-type]
+            for prop_name, prop in write_properties.items():
+                read_prop = self._as_read_properties(prop, all_containers)
+                if isinstance(read_prop, dm.MappedProperty) and read_prop.container not in referenced_containers:
+                    referenced_containers[read_prop.container] = all_containers[read_prop.container]
+                properties[prop_name] = read_prop
+
+            read_view = dm.View(
+                space=view.space,
+                external_id=view.external_id,
+                version=view.version,
+                description=view.description,
+                name=view.name,
+                filter=view.filter,
+                implements=view.implements.copy(),
+                used_for=self._used_for(referenced_containers.values()),
+                writable=self._writable(properties.values(), referenced_containers.values()),
+                properties=properties,
+                is_global=False,
+                last_updated_time=0,
+                created_time=0,
+            )
+            views.append(read_view)
+
+        return dm.DataModel(
+            space=self.data_model.space,
+            external_id=self.data_model.external_id,
+            version=self.data_model.version,
+            name=self.data_model.name,
+            description=self.data_model.description,
+            views=views,
+            is_global=False,
+            last_updated_time=0,
+            created_time=0,
+        )
+
+    @staticmethod
+    def _as_read_properties(
+        write: ViewPropertyApply, all_containers: MutableMapping[dm.ContainerId, dm.ContainerApply]
+    ) -> ViewProperty:
+        if isinstance(write, dm.MappedPropertyApply):
+            container_prop = all_containers[write.container].properties[write.container_property_identifier]
+            return dm.MappedProperty(
+                container=write.container,
+                container_property_identifier=write.container_property_identifier,
+                name=write.name,
+                description=write.description,
+                source=write.source,
+                type=container_prop.type,
+                nullable=container_prop.nullable,
+                auto_increment=container_prop.auto_increment,
+                immutable=container_prop.immutable,
+                # Likely bug in SDK.
+                default_value=container_prop.default_value,  # type: ignore[arg-type]
+            )
+        if isinstance(write, dm.EdgeConnectionApply):
+            edge_cls = SingleEdgeConnection if isinstance(write, SingleEdgeConnectionApply) else dm.MultiEdgeConnection
+            return edge_cls(
+                type=write.type,
+                source=write.source,
+                name=write.name,
+                description=write.description,
+                edge_source=write.edge_source,
+                direction=write.direction,
+            )
+        if isinstance(write, ReverseDirectRelationApply):
+            relation_cls = (
+                SingleReverseDirectRelation
+                if isinstance(write, SingleReverseDirectRelationApply)
+                else dm.MultiReverseDirectRelation
+            )
+            return relation_cls(
+                source=write.source,
+                through=write.through,
+                name=write.name,
+                description=write.description,
+            )
+        raise ValueError(f"Cannot convert {write} to read format")
+
+    @staticmethod
+    def _used_for(containers: Iterable[dm.ContainerApply]) -> Literal["node", "edge", "all"]:
+        used_for = {container.used_for for container in containers}
+        if used_for == {"node"}:
+            return "node"
+        if used_for == {"edge"}:
+            return "edge"
+        return "all"
+
+    @staticmethod
+    def _writable(properties: Iterable[ViewProperty], containers: Iterable[dm.ContainerApply]) -> bool:
+        used_properties = {
+            (prop.container, prop.container_property_identifier)
+            for prop in properties
+            if isinstance(prop, dm.MappedProperty)
+        }
+        required_properties = {
+            (container.as_id(), prop_id)
+            for container in containers
+            for prop_id, prop in container.properties.items()
+            if not prop.nullable
+        }
+        # If a container has a required property that is not used by the view, the view is not writable
+        return not bool(required_properties - used_properties)

cognite/neat/_constants.py

@@ -1,7 +1,7 @@
 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 rdflib import DC, DCTERMS, FOAF, OWL, RDF, RDFS, SH, SKOS, XSD, Namespace, URIRef
 
 from cognite import neat
 
@@ -36,7 +36,26 @@ COGNITE_MODELS = (
     DataModelId("cdf_cdm", "CogniteCore", "v1"),
     DataModelId("cdf_idm", "CogniteProcessIndustries", "v1"),
 )
-COGNITE_SPACES = frozenset(model.space for model in COGNITE_MODELS)
+COGNITE_SPACES = frozenset(
+    {model.space for model in COGNITE_MODELS}
+    | {
+        "cdf_360_image_schema",
+        "cdf_3d_schema",
+        "cdf_apm",
+        "cdf_apps_shared",
+        "cdf_cdm",
+        "cdf_cdm_3d",
+        "cdf_cdm_units",
+        "cdf_classic",
+        "cdf_core",
+        "cdf_extraction_extensions",
+        "cdf_idm",
+        "cdf_industrial_canvas",
+        "cdf_infield",
+        "cdf_time_series_data",
+        "cdf_units",
+    }
+)
 DMS_LISTABLE_PROPERTY_LIMIT = 1000
 
 EXAMPLE_RULES = PACKAGE_DIRECTORY / "_rules" / "examples"
@@ -54,13 +73,17 @@ XML_SCHEMA_NAMESPACE = Namespace("http://www.w3.org/2001/XMLSchema#")
 
 def get_default_prefixes() -> dict[str, Namespace]:
     return {
+        "owl": OWL._NS,
         "rdf": RDF._NS,
         "rdfs": RDFS._NS,
-        "dct": DCTERMS._NS,
+        "dcterms": DCTERMS._NS,
+        "dc": DC._NS,
         "skos": SKOS._NS,
-        "owl": OWL._NS,
+        "sh": SH._NS,
         "xsd": XSD._NS,
+        "imf": Namespace("http://ns.imfid.org/imf#"),
         "pav": Namespace("http://purl.org/pav/"),
+        "foaf": FOAF._NS,
     }
 
 

cognite/neat/_graph/_shared.py

@@ -4,31 +4,30 @@ MIMETypes: TypeAlias = Literal[
     "application/rdf+xml", "text/turtle", "application/n-triple", "application/n-quads", "application/trig"
 ]
 
+RDFTypes: TypeAlias = Literal["xml", "rdf", "owl", "n3", "ttl", "turtle", "nt", "nq", "nquads", "trig"]
 
-def rdflib_to_mime_types(rdflib_format: str) -> str | None:
+
+def rdflib_to_oxi_type(rdflib_format: str) -> str | None:
     """Convert an RDFlib format to a MIME type.
 
     Args:
         rdflib_format (str): The RDFlib format.
 
     Returns:
-        MIMETypes: The MIME type.
-
-    !!! note
-        This will be replaced once new version of oxrdflib is released.
+        Oxi format used to trigger correct plugging in rdflib
 
     """
 
     mapping = {
-        "xml": "application/rdf+xml",
-        "rdf": "application/rdf+xml",
-        "owl": "application/rdf+xml",
-        "n3": "application/n-triple",
-        "ttl": "text/turtle",
-        "turtle": "text/turtle",
-        "nt": "application/n-triple",
-        "nq": "application/n-quads",
-        "nquads": "application/n-quads",
-        "trig": "application/trig",
+        "xml": "ox-xml",
+        "rdf": "ox-xml",
+        "owl": "ox-xml",
+        "n3": "ox-n3",
+        "ttl": "ox-ttl",
+        "turtle": "ox-turtle",
+        "nt": "ox-nt",
+        "nq": "ox-nq",
+        "nquads": "ox-nquads",
+        "trig": "ox-trig",
     }
     return mapping.get(rdflib_format, None)