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

cognite/neat/_session/_collector.py

@@ -0,0 +1,126 @@
+import os
+import platform
+import tempfile
+import threading
+import uuid
+from contextlib import suppress
+from functools import cached_property
+from pathlib import Path
+from typing import Any
+
+from mixpanel import Consumer, Mixpanel  # type: ignore[import-untyped]
+
+from cognite.neat._constants import IN_NOTEBOOK, IN_PYODIDE
+from cognite.neat._version import __version__
+
+_NEAT_MIXPANEL_TOKEN: str = "bd630ad149d19999df3989c3a3750c4f"
+
+
+class Collector:
+    def __init__(self, skip_tracking: bool = False) -> None:
+        self.mp = Mixpanel(_NEAT_MIXPANEL_TOKEN, consumer=Consumer(api_host="api-eu.mixpanel.com"))
+        tmp_dir = Path(tempfile.gettempdir()).resolve()
+        self._opt_status_file = tmp_dir / "neat-opt-status.bin"
+        self._distinct_id_file = tmp_dir / "neat-distinct-id.bin"
+        self.skip_tracking = self.opted_out or skip_tracking
+
+    @cached_property
+    def _opt_status(self) -> str:
+        if self._opt_status_file.exists():
+            return self._opt_status_file.read_text()
+        return ""
+
+    def _bust_opt_status(self) -> None:
+        self.__dict__.pop("_opt_status", None)
+
+    @property
+    def opted_out(self) -> bool:
+        return self._opt_status == "opted-out"
+
+    @property
+    def opted_in(self) -> bool:
+        return self._opt_status == "opted-in"
+
+    @staticmethod
+    def _get_environment() -> str:
+        if IN_PYODIDE:
+            return "pyodide"
+        if IN_NOTEBOOK:
+            return "notebook"
+        return "python"
+
+    def track_session_command(self, command: str, *args, **kwargs) -> None:
+        event_information = {
+            "neatVersion": __version__,
+            "$os": platform.system(),
+            "pythonVersion": platform.python_version(),
+            "environment": self._get_environment(),
+        }
+
+        if len(args) > 1:
+            # The first argument is self.
+            for i, arg in enumerate(args[1:]):
+                event_information[f"arg{i}"] = arg
+
+        if kwargs:
+            for key, value in kwargs.items():
+                event_information[key] = self._serialize_value(value)[:500]
+        self._track(command, event_information)
+
+    @staticmethod
+    def _serialize_value(value: Any) -> str:
+        if isinstance(value, (str | int | float | bool)):
+            return str(value)
+        if isinstance(value, list | tuple | dict):
+            return str(value)
+        return str(type(value))
+
+    def _track(self, event_name: str, event_information: dict[str, Any]) -> bool:
+        if self.skip_tracking or not self.opted_in or "PYTEST_CURRENT_TEST" in os.environ:
+            return False
+
+        distinct_id = self.get_distinct_id()
+
+        def track() -> None:
+            # If we are unable to connect to Mixpanel, we don't want to crash the program
+            with suppress(ConnectionError):
+                self.mp.track(
+                    distinct_id,
+                    event_name,
+                    event_information,
+                )
+
+        thread = threading.Thread(
+            target=track,
+            daemon=False,
+        )
+        thread.start()
+        return True
+
+    def get_distinct_id(self) -> str:
+        if self._distinct_id_file.exists():
+            return self._distinct_id_file.read_text()
+
+        distinct_id = f"{platform.system()}-{platform.python_version()}-{uuid.uuid4()!s}"
+        self._distinct_id_file.write_text(distinct_id)
+        with suppress(ConnectionError):
+            self.mp.people_set(
+                distinct_id,
+                {
+                    "$os": platform.system(),
+                    "$python_version": platform.python_version(),
+                    "$distinct_id": distinct_id,
+                },
+            )
+        return distinct_id
+
+    def enable(self) -> None:
+        self._opt_status_file.write_text("opted-in")
+        self._bust_opt_status()
+
+    def disable(self) -> None:
+        self._opt_status_file.write_text("opted-out")
+        self._bust_opt_status()
+
+
+_COLLECTOR = Collector()

cognite/neat/_session/_inspect.py

@@ -9,10 +9,10 @@ from cognite.neat._issues import IssueList
 from cognite.neat._utils.upload import UploadResult, UploadResultCore, UploadResultList
 
 from ._state import SessionState
-from .exceptions import intercept_session_exceptions
+from .exceptions import session_class_wrapper
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class InspectAPI:
     def __init__(self, state: SessionState) -> None:
         self._state = state
@@ -25,7 +25,7 @@ class InspectAPI:
         return self._state.data_model.last_verified_rule[1].properties.to_pandas()
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class InspectIssues:
     """Inspect issues of the current data model."""
 
@@ -92,14 +92,14 @@ class InspectIssues:
         )
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class InspectOutcome:
     def __init__(self, state: SessionState) -> None:
         self.data_model = InspectUploadOutcome(lambda: state.data_model.last_outcome)
         self.instances = InspectUploadOutcome(lambda: state.instances.last_outcome)
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class InspectUploadOutcome:
     def __init__(self, get_last_outcome: Callable[[], UploadResultList]) -> None:
         self._get_last_outcome = get_last_outcome

cognite/neat/_session/_prepare.py

@@ -12,10 +12,10 @@ from cognite.neat._rules.transformers import ReduceCogniteModel, ToCompliantEnti
 from cognite.neat._store._provenance import Change
 
 from ._state import SessionState
-from .exceptions import NeatSessionError, intercept_session_exceptions
+from .exceptions import NeatSessionError, session_class_wrapper
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class PrepareAPI:
     def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state
@@ -24,7 +24,7 @@ class PrepareAPI:
         self.instances = InstancePrepareAPI(state, verbose)
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class InstancePrepareAPI:
     def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state
@@ -95,7 +95,7 @@ class InstancePrepareAPI:
         return type_uri[0], property_uri[0]
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class DataModelPrepareAPI:
     def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state

cognite/neat/_session/_read.py

@@ -1,17 +1,19 @@
 import tempfile
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes.data_modeling import DataModelId, DataModelIdentifier
 
+from cognite.neat._constants import COGNITE_SPACES
 from cognite.neat._graph import examples as instances_examples
 from cognite.neat._graph import extractors
 from cognite.neat._issues import IssueList
 from cognite.neat._issues.errors import NeatValueError
 from cognite.neat._rules import importers
 from cognite.neat._rules._shared import ReadRules
+from cognite.neat._rules.importers import BaseImporter
 from cognite.neat._store._provenance import Activity as ProvenanceActivity
 from cognite.neat._store._provenance import Change
 from cognite.neat._store._provenance import Entity as ProvenanceEntity
@@ -20,10 +22,10 @@ from cognite.neat._utils.reader import GitHubReader, NeatReader, PathReader
 from ._state import SessionState
 from ._wizard import NeatObjectType, RDFFileType, object_wizard, rdf_dm_wizard
 from .engine import import_engine
-from .exceptions import NeatSessionError, intercept_session_exceptions
+from .exceptions import NeatSessionError, session_class_wrapper
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class ReadAPI:
     def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
         self._state = state
@@ -32,9 +34,10 @@ class ReadAPI:
         self.rdf = RDFReadAPI(state, client, verbose)
         self.excel = ExcelReadAPI(state, client, verbose)
         self.csv = CSVReadAPI(state, client, verbose)
+        self.yaml = YamlReadAPI(state, client, verbose)
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class BaseReadAPI:
     def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
         self._state = state
@@ -62,7 +65,7 @@ class BaseReadAPI:
             raise NeatValueError(f"Expected str or Path, got {type(io)}")
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class CDFReadAPI(BaseReadAPI):
     def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
         super().__init__(state, client, verbose)
@@ -107,7 +110,7 @@ class CDFReadAPI(BaseReadAPI):
         return self._store_rules(rules, change)
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class CDFClassicAPI(BaseReadAPI):
     @property
     def _get_client(self) -> CogniteClient:
@@ -122,7 +125,7 @@ class CDFClassicAPI(BaseReadAPI):
             print(f"Asset hierarchy {root_asset_external_id} read successfully")
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class ExcelReadAPI(BaseReadAPI):
     def __call__(self, io: Any) -> IssueList:
         reader = NeatReader.create(io)
@@ -146,7 +149,57 @@ class ExcelReadAPI(BaseReadAPI):
         return input_rules.issues
 
 
-@intercept_session_exceptions
+@session_class_wrapper
+class YamlReadAPI(BaseReadAPI):
+    def __call__(self, io: Any, format: Literal["neat", "toolkit"] = "neat") -> IssueList:
+        reader = NeatReader.create(io)
+        if not isinstance(reader, PathReader):
+            raise NeatValueError("Only file paths are supported for YAML files")
+        start = datetime.now(timezone.utc)
+        importer: BaseImporter
+        if format == "neat":
+            importer = importers.YAMLImporter.from_file(reader.path)
+        elif format == "toolkit":
+            if reader.path.is_file():
+                dms_importer = importers.DMSImporter.from_zip_file(reader.path)
+            elif reader.path.is_dir():
+                dms_importer = importers.DMSImporter.from_directory(reader.path)
+            else:
+                raise NeatValueError(f"Unsupported YAML format: {format}")
+            ref_containers = dms_importer.root_schema.referenced_container()
+            if system_container_ids := [
+                container_id for container_id in ref_containers if container_id.space in COGNITE_SPACES
+            ]:
+                if self._client is None:
+                    raise NeatSessionError(
+                        "No client provided. You are referencing Cognite containers in your data model, "
+                        "NEAT needs a client to lookup the container definitions. "
+                        "Please set the client in the session, NeatSession(client=client)."
+                    )
+                system_containers = self._client.data_modeling.containers.retrieve(system_container_ids)
+                dms_importer.update_referenced_containers(system_containers)
+
+            importer = dms_importer
+        else:
+            raise NeatValueError(f"Unsupported YAML format: {format}")
+        input_rules: ReadRules = importer.to_rules()
+
+        end = datetime.now(timezone.utc)
+
+        if input_rules.rules:
+            change = Change.from_rules_activity(
+                input_rules,
+                importer.agent,
+                start,
+                end,
+                description=f"YAML file {reader!s} read as unverified data model",
+            )
+            self._store_rules(input_rules, change)
+
+        return input_rules.issues
+
+
+@session_class_wrapper
 class CSVReadAPI(BaseReadAPI):
     def __call__(self, io: Any, type: str, primary_key: str) -> None:
         reader = NeatReader.create(io)
@@ -167,7 +220,7 @@ class CSVReadAPI(BaseReadAPI):
         self._state.instances.store.write(extractor)
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class RDFReadAPI(BaseReadAPI):
     def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
         super().__init__(state, client, verbose)

cognite/neat/_session/_set.py

@@ -7,10 +7,10 @@ from cognite.neat._rules.transformers import SetIDDMSModel
 from cognite.neat._store._provenance import Change
 
 from ._state import SessionState
-from .exceptions import NeatSessionError, intercept_session_exceptions
+from .exceptions import NeatSessionError, session_class_wrapper
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class SetAPI:
     def __init__(self, state: SessionState, verbose: bool) -> None:
         self._state = state

cognite/neat/_session/_show.py

@@ -16,10 +16,10 @@ from cognite.neat._session.exceptions import NeatSessionError
 from cognite.neat._utils.rdf_ import remove_namespace_from_uri
 
 from ._state import SessionState
-from .exceptions import intercept_session_exceptions
+from .exceptions import session_class_wrapper
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class ShowAPI:
     def __init__(self, state: SessionState) -> None:
         self._state = state
@@ -27,7 +27,7 @@ class ShowAPI:
         self.instances = ShowInstanceAPI(self._state)
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class ShowBaseAPI:
     def __init__(self, state: SessionState) -> None:
         self._state = state
@@ -63,7 +63,7 @@ class ShowBaseAPI:
             return net.show(name)
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class ShowDataModelAPI(ShowBaseAPI):
     def __init__(self, state: SessionState) -> None:
         super().__init__(state)
@@ -111,7 +111,7 @@ class ShowDataModelAPI(ShowBaseAPI):
                 di_graph.add_edge(
                     prop_.view.suffix,
                     prop_.value_type.suffix,
-                    label=prop_.name or prop_.property_,
+                    label=prop_.name or prop_.view_property,
                 )
 
         return di_graph
@@ -145,7 +145,7 @@ class ShowDataModelAPI(ShowBaseAPI):
         return di_graph
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class ShowDataModelImplementsAPI(ShowBaseAPI):
     def __init__(self, state: SessionState) -> None:
         super().__init__(state)
@@ -206,27 +206,27 @@ class ShowDataModelImplementsAPI(ShowBaseAPI):
             # if possible use human readable label coming from the view name
 
             # add subClassOff as edges
-            if class_.parent:
+            if class_.implements:
                 if not di_graph.has_node(class_.class_.suffix):
                     di_graph.add_node(
                         class_.class_.suffix,
                         label=class_.name or class_.class_.suffix,
                     )
 
-                for parent in class_.parent:
+                for parent in class_.implements:
                     if not di_graph.has_node(parent.suffix):
                         di_graph.add_node(parent.suffix, label=parent.suffix)
                     di_graph.add_edge(
                         class_.class_.suffix,
                         parent.suffix,
-                        label="subClassOf",
+                        label="implements",
                         dashes=True,
                     )
 
         return di_graph
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class ShowDataModelProvenanceAPI(ShowBaseAPI):
     def __init__(self, state: SessionState) -> None:
         super().__init__(state)
@@ -286,7 +286,7 @@ class ShowDataModelProvenanceAPI(ShowBaseAPI):
         return remove_namespace_from_uri(thing)
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class ShowInstanceAPI(ShowBaseAPI):
     def __init__(self, state: SessionState) -> None:
         super().__init__(state)

cognite/neat/_session/_to.py

@@ -12,10 +12,10 @@ from cognite.neat._rules._shared import VerifiedRules
 from cognite.neat._utils.upload import UploadResultCore, UploadResultList
 
 from ._state import SessionState
-from .exceptions import NeatSessionError, intercept_session_exceptions
+from .exceptions import NeatSessionError, session_class_wrapper
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class ToAPI:
     def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
         self._state = state
@@ -46,21 +46,34 @@ class ToAPI:
         return None
 
     @overload
-    def yaml(self, io: None) -> str: ...
+    def yaml(self, io: None, format: Literal["neat"] = "neat") -> str: ...
 
     @overload
-    def yaml(self, io: Any) -> None: ...
-
-    def yaml(self, io: Any | None = None) -> str | None:
-        exporter = exporters.YAMLExporter()
-        if io is None:
-            return exporter.export(self._state.data_model.last_verified_rule[1])
+    def yaml(self, io: Any, format: Literal["neat", "toolkit"] = "neat") -> None: ...
+
+    def yaml(self, io: Any | None = None, format: Literal["neat", "toolkit"] = "neat") -> str | None:
+        if format == "neat":
+            exporter = exporters.YAMLExporter()
+            last_verified = self._state.data_model.last_verified_rule[1]
+            if io is None:
+                return exporter.export(last_verified)
+
+            exporter.export_to_file(last_verified, Path(io))
+        elif format == "toolkit":
+            if io is None or not isinstance(io, str | Path):
+                raise NeatSessionError(
+                    "Please provide a zip file or directory path to write the YAML files to."
+                    "This is required for the 'toolkit' format."
+                )
+            dms_rule = self._state.data_model.last_verified_dms_rules[1]
+            exporters.DMSExporter().export_to_file(dms_rule, Path(io))
+        else:
+            raise NeatSessionError("Please provide a valid format. {['neat', 'toolkit']}")
 
-        exporter.export_to_file(self._state.data_model.last_verified_rule[1], Path(io))
         return None
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class CDFToAPI:
     def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
         self._client = client

cognite/neat/_session/exceptions.py

@@ -2,6 +2,8 @@ import functools
 from collections.abc import Callable
 from typing import Any
 
+from ._collector import _COLLECTOR
+
 try:
     from rich import print
 
@@ -16,9 +18,10 @@ class NeatSessionError(Exception):
     ...
 
 
-def _intercept_session_exceptions(func: Callable):
+def _session_method_wrapper(func: Callable, cls_name: str):
     @functools.wraps(func)
     def wrapper(*args: Any, **kwargs: Any):
+        _COLLECTOR.track_session_command(f"{cls_name}.{func.__name__}", *args, **kwargs)
         try:
             return func(*args, **kwargs)
         except NeatSessionError as e:
@@ -40,7 +43,21 @@ def _intercept_session_exceptions(func: Callable):
     return wrapper
 
 
-def intercept_session_exceptions(cls: type):
+def session_class_wrapper(cls: type):
+    """This decorator wraps all methods of a class.
+
+    It should be used with all composition classes used with the NeatSession class.
+
+    It does the following:
+        * Intercepts all NeatSession exceptions and prints them in a user-friendly way.
+        * Collects user metrics.
+
+    Args:
+        cls: NeatSession composition class
+
+    Returns:
+        cls: NeatSession composition class with all methods wrapped
+    """
     to_check = [cls]
     while to_check:
         cls = to_check.pop()
@@ -49,7 +66,7 @@ def intercept_session_exceptions(cls: type):
                 continue
             attr = getattr(cls, attr_name)
             if callable(attr):
-                setattr(cls, attr_name, _intercept_session_exceptions(attr))
+                setattr(cls, attr_name, _session_method_wrapper(attr, cls.__name__))
             elif isinstance(attr, type):
                 to_check.append(attr)
     return cls

cognite/neat/_store/_provenance.py

@@ -93,14 +93,14 @@ class Entity:
             return cls(
                 was_attributed_to=agent,
                 was_generated_by=activity,
-                id_=rules.id_,
+                id_=rules.metadata.identifier,
             )
 
         elif isinstance(rules, ReadRules | JustRules) and rules.rules is not None:
             return cls(
                 was_attributed_to=agent,
                 was_generated_by=activity,
-                id_=rules.rules.id_,
+                id_=rules.rules.metadata.identifier,
             )
         else:
             return cls(

cognite/neat/_utils/auxiliary.py

@@ -140,3 +140,22 @@ def string_to_ideal_type(input_string: str) -> int | bool | float | datetime | s
                 except ValueError:
                     # Return the input string if no conversion is possible
                     return input_string
+
+
+def get_parameters_by_method(obj: object, prefix: str = "") -> dict[str, dict[str, type]]:
+    annotation_by_method = {}
+    namespace_dict = {**dict(obj.__class__.__dict__.items()), **vars(obj)}
+    for name, value in namespace_dict.items():
+        if name != "__call__" and (name.startswith("_") or isinstance(value, property)):
+            continue
+        if callable(value) and type(value).__name__ == "function":
+            annotation_by_method[f"{prefix}{name}"] = get_parameters(value)
+        elif isinstance(value, object) and type(value).__module__ != "builtins":
+            sub_prefix = f"{prefix}{name}." if prefix else f"{name}."
+            annotation_by_method.update(get_parameters_by_method(value, sub_prefix))
+    return annotation_by_method
+
+
+def get_parameters(obj: Callable) -> dict[str, type]:
+    annotations = inspect.get_annotations(obj)
+    return {name: annotations[name] for name in annotations if name != "return"}

cognite/neat/_version.py

@@ -1,2 +1,2 @@
-__version__ = "0.97.3"
+__version__ = "0.98.0"
 __engine__ = "^1.0.3"

cognite/neat/_workflows/steps/data_contracts.py

@@ -10,9 +10,7 @@ from cognite.client.data_classes import (
 from cognite.client.data_classes.data_modeling import EdgeApply, NodeApply
 
 from cognite.neat._rules.models import (
-    AssetRules,
     DMSRules,
-    DomainRules,
     InformationRules,
 )
 from cognite.neat._store import NeatGraphStore
@@ -20,19 +18,13 @@ from cognite.neat._workflows.steps.step_model import DataContract
 
 
 class MultiRuleData(DataContract):
-    domain: DomainRules | None = None
     information: InformationRules | None = None
-    asset: AssetRules | None = None
     dms: DMSRules | None = None
 
     @classmethod
-    def from_rules(cls, rules: DomainRules | InformationRules | AssetRules | DMSRules):
-        if isinstance(rules, DomainRules):
-            return cls(domain=rules)
-        elif isinstance(rules, InformationRules):
+    def from_rules(cls, rules: InformationRules | DMSRules):
+        if isinstance(rules, InformationRules):
             return cls(information=rules)
-        elif isinstance(rules, AssetRules):
-            return cls(asset=rules)
         elif isinstance(rules, DMSRules):
             return cls(dms=rules)
         else: