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

cognite/neat/_rules/transformers/_mapping.py

@@ -1,13 +1,16 @@
+import warnings
 from abc import ABC
 from collections import defaultdict
+from functools import cached_property
+from typing import Any, ClassVar, Literal
 
+from cognite.neat._issues.errors import NeatValueError
+from cognite.neat._issues.warnings import NeatValueWarning, PropertyOverwritingWarning
 from cognite.neat._rules._shared import JustRules, OutRules
-from cognite.neat._rules.models import DMSRules, InformationRules
-from cognite.neat._rules.models._base_rules import ClassRef
-from cognite.neat._rules.models.dms import DMSProperty
-from cognite.neat._rules.models.entities import ClassEntity, ReferenceEntity
-from cognite.neat._rules.models.information import InformationClass
-from cognite.neat._rules.models.mapping import RuleMapping
+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 ._base import RulesTransformer
 
@@ -51,9 +54,6 @@ class MapOneToOne(MapOntoTransformers):
 
     def transform(self, rules: DMSRules | OutRules[DMSRules]) -> JustRules[DMSRules]:
         solution: DMSRules = self._to_rules(rules)
-        if solution.reference is not None:
-            raise ValueError("Reference already exists")
-        solution.reference = self.reference
         view_by_external_id = {view.view.external_id: view for view in solution.views}
         ref_view_by_external_id = {view.view.external_id: view for view in self.reference.views}
 
@@ -95,12 +95,11 @@ class MapOneToOne(MapOntoTransformers):
                 if ref_prop.container and ref_prop.container_property:
                     prop.container = ref_prop.container
                     prop.container_property = ref_prop.container_property
-                prop.reference = ReferenceEntity.from_entity(ref_prop.view, ref_prop.view_property)
 
         return JustRules(solution)
 
 
-class RuleMapper(RulesTransformer[InformationRules, InformationRules]):
+class RuleMapper(RulesTransformer[DMSRules, DMSRules]):
     """Maps properties and classes using the given mapping.
 
     **Note**: This transformer mutates the input rules.
@@ -110,30 +109,102 @@ class RuleMapper(RulesTransformer[InformationRules, InformationRules]):
 
     """
 
-    def __init__(self, mapping: RuleMapping) -> None:
+    _mapping_fields: ClassVar[frozenset[str]] = frozenset(
+        ["connection", "value_type", "nullable", "immutable", "is_list", "default", "index", "constraint"]
+    )
+
+    def __init__(self, mapping: DMSRules, data_type_conflict: Literal["overwrite"] = "overwrite") -> None:
         self.mapping = mapping
+        self.data_type_conflict = data_type_conflict
+
+    @cached_property
+    def _view_by_entity_id(self) -> dict[str, DMSView]:
+        return {view.view.external_id: view for view in self.mapping.views}
 
-    def transform(self, rules: InformationRules | OutRules[InformationRules]) -> JustRules[InformationRules]:
+    @cached_property
+    def _property_by_view_property(self) -> dict[tuple[str, str], DMSProperty]:
+        return {(prop.view.external_id, prop.view_property): prop for prop in self.mapping.properties}
+
+    def transform(self, rules: DMSRules | OutRules[DMSRules]) -> JustRules[DMSRules]:
+        if self.data_type_conflict != "overwrite":
+            raise NeatValueError(f"Invalid data_type_conflict: {self.data_type_conflict}")
         input_rules = self._to_rules(rules)
+        new_rules = input_rules.model_copy(deep=True)
+        new_rules.metadata.version += "_mapped"
 
-        destination_by_source = self.mapping.properties.as_destination_by_source()
-        destination_cls_by_source = self.mapping.classes.as_destination_by_source()
-        used_destination_classes: set[ClassEntity] = set()
-        for prop in input_rules.properties:
-            if destination_prop := destination_by_source.get(prop.as_reference()):
-                prop.class_ = destination_prop.class_
-                prop.property_ = destination_prop.property_
-                used_destination_classes.add(destination_prop.class_)
-            elif destination_cls := destination_cls_by_source.get(ClassRef(Class=prop.class_)):
-                # If the property is not in the mapping, but the class is,
-                # then we should map the class to the destination
-                prop.class_ = destination_cls.class_
-
-        for cls_ in input_rules.classes:
-            if destination_cls := destination_cls_by_source.get(cls_.as_reference()):
-                cls_.class_ = destination_cls.class_
-        existing_classes = {cls_.class_ for cls_ in input_rules.classes}
-        for new_cls in used_destination_classes - existing_classes:
-            input_rules.classes.append(InformationClass(class_=new_cls))
-
-        return JustRules(input_rules)
+        for view in new_rules.views:
+            if mapping_view := self._view_by_entity_id.get(view.view.external_id):
+                view.implements = mapping_view.implements
+
+        for prop in new_rules.properties:
+            key = (prop.view.external_id, prop.view_property)
+            if key not in self._property_by_view_property:
+                continue
+            mapping_prop = self._property_by_view_property[key]
+            to_overwrite, conflicts = self._find_overwrites(prop, mapping_prop)
+            if conflicts and self.data_type_conflict == "overwrite":
+                warnings.warn(
+                    PropertyOverwritingWarning(prop.view.as_id(), "view", prop.view_property, tuple(conflicts)),
+                    stacklevel=2,
+                )
+            elif conflicts:
+                raise NeatValueError(f"Conflicting properties for {prop.view}.{prop.view_property}: {conflicts}")
+            for field_name, value in to_overwrite.items():
+                setattr(prop, field_name, value)
+            prop.container = mapping_prop.container
+            prop.container_property = mapping_prop.container_property
+
+        # Add missing views used as value types
+        existing_views = {view.view for view in new_rules.views}
+        new_value_types = {
+            prop.value_type
+            for prop in new_rules.properties
+            if isinstance(prop.value_type, ViewEntity) and prop.value_type not in existing_views
+        }
+        for new_value_type in new_value_types:
+            if mapping_view := self._view_by_entity_id.get(new_value_type.external_id):
+                new_rules.views.append(mapping_view)
+            else:
+                warnings.warn(NeatValueWarning(f"View {new_value_type} not found in mapping"), stacklevel=2)
+
+        # Add missing enums
+        existing_enum_collections = {item.collection for item in new_rules.enum or []}
+        new_enums = {
+            prop.value_type.collection
+            for prop in new_rules.properties
+            if isinstance(prop.value_type, Enum) and prop.value_type.collection not in existing_enum_collections
+        }
+        if new_enums:
+            new_rules.enum = new_rules.enum or SheetList[DMSEnum]([])
+            for item in self.mapping.enum or []:
+                if item.collection in new_enums:
+                    new_rules.enum.append(item)
+
+        return JustRules(new_rules)
+
+    def _find_overwrites(self, prop: DMSProperty, mapping_prop: DMSProperty) -> tuple[dict[str, Any], list[str]]:
+        """Finds the properties that need to be overwritten and returns them.
+
+        In addition, conflicting properties are returned. Note that overwriting properties that are
+        originally None is not considered a conflict. Thus, you can have properties to overwrite but no
+        conflicts.
+
+        Args:
+            prop: The property to compare.
+            mapping_prop: The property to compare against.
+
+        Returns:
+            A tuple with the properties to overwrite and the conflicting properties.
+
+        """
+        to_overwrite: dict[str, Any] = {}
+        conflicts: list[str] = []
+        for field_name in self._mapping_fields:
+            mapping_value = getattr(mapping_prop, field_name)
+            source_value = getattr(prop, field_name)
+            if mapping_value != source_value:
+                to_overwrite[field_name] = mapping_value
+                if source_value is not None:
+                    # 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

cognite/neat/_rules/transformers/_verification.py

@@ -13,12 +13,8 @@ from cognite.neat._rules._shared import (
     VerifiedRules,
 )
 from cognite.neat._rules.models import (
-    AssetInputRules,
-    AssetRules,
     DMSInputRules,
     DMSRules,
-    DomainInputRules,
-    DomainRules,
     InformationInputRules,
     InformationRules,
 )
@@ -31,8 +27,9 @@ class VerificationTransformer(RulesTransformer[T_InputRules, T_VerifiedRules], A
 
     _rules_cls: type[T_VerifiedRules]
 
-    def __init__(self, errors: Literal["raise", "continue"]) -> None:
+    def __init__(self, errors: Literal["raise", "continue"], post_validate: bool = True) -> None:
         self.errors = errors
+        self.post_validate = post_validate
 
     def transform(self, rules: T_InputRules | OutRules[T_InputRules]) -> MaybeRules[T_VerifiedRules]:
         issues = IssueList()
@@ -43,7 +40,9 @@ class VerificationTransformer(RulesTransformer[T_InputRules, T_VerifiedRules], A
         verified_rules: T_VerifiedRules | None = None
         with catch_issues(issues, NeatError, NeatWarning, error_args) as future:
             rules_cls = self._get_rules_cls(in_)
-            verified_rules = rules_cls.model_validate(in_.dump())  # type: ignore[assignment]
+            dumped = in_.dump()
+            dumped["post_validate"] = self.post_validate
+            verified_rules = rules_cls.model_validate(dumped)  # type: ignore[assignment]
 
         if (future.result == "failure" or issues.has_errors or verified_rules is None) and self.errors == "raise":
             raise issues.as_errors()
@@ -68,12 +67,6 @@ class VerifyInformationRules(VerificationTransformer[InformationInputRules, Info
     _rules_cls = InformationRules
 
 
-class VerifyAssetRules(VerificationTransformer[AssetInputRules, AssetRules]):
-    """Class to verify Asset rules."""
-
-    _rules_cls = AssetRules
-
-
 class VerifyAnyRules(VerificationTransformer[InputRules, VerifiedRules]):
     """Class to verify arbitrary rules"""
 
@@ -82,9 +75,5 @@ class VerifyAnyRules(VerificationTransformer[InputRules, VerifiedRules]):
             return InformationRules
         elif isinstance(in_, DMSInputRules):
             return DMSRules
-        elif isinstance(in_, AssetInputRules):
-            return AssetRules
-        elif isinstance(in_, DomainInputRules):
-            return DomainRules
         else:
             raise NeatTypeError(f"Unsupported rules type: {type(in_)}")

cognite/neat/_session/_base.py

@@ -5,14 +5,13 @@ from cognite.client import CogniteClient
 from cognite.client import data_modeling as dm
 
 from cognite.neat import _version
+from cognite.neat._client import NeatClient
 from cognite.neat._issues import IssueList, catch_issues
 from cognite.neat._issues.errors import RegexViolationError
 from cognite.neat._rules import importers
 from cognite.neat._rules._shared import ReadRules, VerifiedRules
-from cognite.neat._rules.importers._rdf._base import DEFAULT_NON_EXISTING_NODE_TYPE
-from cognite.neat._rules.models import DMSRules
-from cognite.neat._rules.models.data_types import AnyURI
-from cognite.neat._rules.models.entities._single_value import UnknownEntity
+from cognite.neat._rules.importers import DMSImporter
+from cognite.neat._rules.models import DMSInputRules, DMSRules, SheetList
 from cognite.neat._rules.models.information._rules import InformationRules
 from cognite.neat._rules.models.information._rules_input import InformationInputRules
 from cognite.neat._rules.transformers import ConvertToRules, VerifyAnyRules
@@ -21,9 +20,11 @@ from cognite.neat._store._provenance import (
     INSTANCES_ENTITY,
     Change,
 )
-from cognite.neat._utils.auth import _CLIENT_NAME
 
+from ._collector import _COLLECTOR, Collector
+from ._drop import DropAPI
 from ._inspect import InspectAPI
+from ._mapping import MappingAPI
 from ._prepare import PrepareAPI
 from ._read import ReadAPI
 from ._set import SetAPI
@@ -31,10 +32,10 @@ from ._show import ShowAPI
 from ._state import SessionState
 from ._to import ToAPI
 from .engine import load_neat_engine
-from .exceptions import NeatSessionError, intercept_session_exceptions
+from .exceptions import NeatSessionError, session_class_wrapper
 
 
-@intercept_session_exceptions
+@session_class_wrapper
 class NeatSession:
     def __init__(
         self,
@@ -43,17 +44,19 @@ class NeatSession:
         verbose: bool = True,
         load_engine: Literal["newest", "cache", "skip"] = "cache",
     ) -> None:
-        self._client = client
+        self._client = NeatClient(client) if client else None
         self._verbose = verbose
         self._state = SessionState(store_type=storage)
-        self.read = ReadAPI(self._state, client, verbose)
-        self.to = ToAPI(self._state, client, verbose)
-        self.prepare = PrepareAPI(self._state, verbose)
+        self.read = ReadAPI(self._state, self._client, verbose)
+        self.to = ToAPI(self._state, self._client, verbose)
+        self.prepare = PrepareAPI(self._client, self._state, verbose)
         self.show = ShowAPI(self._state)
         self.set = SetAPI(self._state, verbose)
         self.inspect = InspectAPI(self._state)
-        if self._client is not None and self._client._config is not None:
-            self._client._config.client_name = _CLIENT_NAME
+        self.mapping = MappingAPI(self._state)
+        self.drop = DropAPI(self._state)
+        self.opt = OptAPI()
+        self.opt._display()
         if load_engine != "skip" and (engine_version := load_neat_engine(client, load_engine)):
             print(f"Neat Engine {engine_version} loaded.")
 
@@ -63,6 +66,28 @@ class NeatSession:
 
     def verify(self) -> IssueList:
         source_id, last_unverified_rule = self._state.data_model.last_unverified_rule
+
+        reference_rules: DMSInputRules | None = None
+        if isinstance(last_unverified_rule.rules, DMSInputRules):
+            dms_rules = last_unverified_rule.rules
+            views_ids, containers_ids = dms_rules.imported_views_and_containers_ids()
+            if views_ids or containers_ids:
+                if self._client is None:
+                    raise NeatSessionError(
+                        "No client provided. You are referencing unknown views and containers in your data model, "
+                        "NEAT needs a client to lookup the definitions. "
+                        "Please set the client in the session, NeatSession(client=client)."
+                    )
+                schema = self._client.schema.retrieve(list(views_ids), list(containers_ids))
+
+                importer = DMSImporter(schema)
+                reference_rules = importer.to_rules().rules
+
+                if reference_rules is not None:
+                    dms_rules.views.extend(reference_rules.views)
+                    if dms_rules.containers:
+                        dms_rules.containers.extend(reference_rules.containers or [])
+
         transformer = VerifyAnyRules("continue")
         start = datetime.now(timezone.utc)
         output = transformer.try_transform(last_unverified_rule)
@@ -74,10 +99,26 @@ class NeatSession:
                 transformer.agent,
                 start,
                 end,
-                f"Verified data model {source_id} as {output.rules.id_}",
+                f"Verified data model {source_id} as {output.rules.metadata.identifier}",
                 self._state.data_model.provenance.source_entity(source_id)
                 or self._state.data_model.provenance.target_entity(source_id),
             )
+            if reference_rules is not None and isinstance(output.rules, DMSRules):
+                # Remove the referenced views and containers from the rules
+                ref_view_ids = set(reference_rules.as_view_entities())
+                if ref_view_ids:
+                    output.rules.views = SheetList(
+                        [view for view in output.rules.views if view.view not in ref_view_ids]
+                    )
+                ref_container_ids = reference_rules.as_container_entities()
+                if output.rules.containers and ref_container_ids:
+                    output.rules.containers = SheetList(
+                        [
+                            container
+                            for container in output.rules.containers
+                            if container.container not in ref_container_ids
+                        ]
+                    )
 
             self._state.data_model.write(output.rules, change)
 
@@ -119,7 +160,7 @@ class NeatSession:
                 converter.agent,
                 start,
                 end,
-                f"Converted data model {source_id} to {converted_rules.id_}",
+                f"Converted data model {source_id} to {converted_rules.metadata.identifier}",
                 self._state.data_model.provenance.source_entity(source_id)
                 or self._state.data_model.provenance.target_entity(source_id),
             )
@@ -144,31 +185,28 @@ class NeatSession:
             "NeatInferredDataModel",
             "v1",
         ),
-        non_existing_node_type: UnknownEntity | AnyURI = DEFAULT_NON_EXISTING_NODE_TYPE,
         max_number_of_instance: int = 100,
     ) -> IssueList:
         """Data model inference from instances.
 
         Args:
             model_id: The ID of the inferred data model.
-            non_existing_node_type: The type of node to use when type of node is not possible to determine.
+            max_number_of_instance: The maximum number of instances to use for inference.
         """
-
         model_id = dm.DataModelId.load(model_id)
 
         start = datetime.now(timezone.utc)
         importer = importers.InferenceImporter.from_graph_store(
             store=self._state.instances.store,
-            non_existing_node_type=non_existing_node_type,
             max_number_of_instance=max_number_of_instance,
         )
         inferred_rules: ReadRules = importer.to_rules()
         end = datetime.now(timezone.utc)
 
         if model_id.space:
-            cast(InformationInputRules, inferred_rules.rules).metadata.prefix = model_id.space
+            cast(InformationInputRules, inferred_rules.rules).metadata.space = model_id.space
         if model_id.external_id:
-            cast(InformationInputRules, inferred_rules.rules).metadata.name = model_id.external_id
+            cast(InformationInputRules, inferred_rules.rules).metadata.external_id = model_id.external_id
 
         if model_id.version:
             cast(InformationInputRules, inferred_rules.rules).metadata.version = model_id.version
@@ -208,3 +246,27 @@ class NeatSession:
             output.append(f"<H2>Instances</H2> {state.instances.store._repr_html_()}")
 
         return "<br />".join(output)
+
+
+@session_class_wrapper
+class OptAPI:
+    def __init__(self, collector: Collector | None = None) -> None:
+        self._collector = collector or _COLLECTOR
+
+    def _display(self) -> None:
+        if self._collector.opted_in or self._collector.opted_out:
+            return
+        print(
+            "For Neat to improve, we need to collect usage information. "
+            "You acknowledge and agree that neat may collect usage information."
+            "To remove this message run 'neat.opt.in_() "
+            "or to stop collecting usage information run 'neat.opt.out()'."
+        )
+
+    def in_(self) -> None:
+        self._collector.enable()
+        print("You have successfully opted in to data collection.")
+
+    def out(self) -> None:
+        self._collector.disable()
+        print("You have successfully opted out of data collection.")

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/_drop.py

@@ -0,0 +1,35 @@
+from rdflib import URIRef
+
+from ._state import SessionState
+from .exceptions import session_class_wrapper
+
+try:
+    from rich import print
+except ImportError:
+    ...
+
+
+@session_class_wrapper
+class DropAPI:
+    def __init__(self, state: SessionState):
+        self._state = state
+
+    def instances(self, type: str | list[str]) -> None:
+        """Drop instances from the session.
+
+        Args:
+            type: The type of instances to drop.
+        """
+        type_list = type if isinstance(type, list) else [type]
+        uri_type_type = dict((v, k) for k, v in self._state.instances.store.queries.types.items())
+        selected_uri_by_type: dict[URIRef, str] = {}
+        for type_item in type_list:
+            if type_item not in uri_type_type:
+                print(f"Type {type_item} not found.")
+            selected_uri_by_type[uri_type_type[type_item]] = type_item
+
+        result = self._state.instances.store.queries.drop_types(list(selected_uri_by_type.keys()))
+
+        for type_uri, count in result.items():
+            print(f"Dropped {count} instances of type {selected_uri_by_type[type_uri]}")
+        return None

cognite/neat/_session/_inspect.py

@@ -9,10 +9,17 @@ 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
 
+try:
+    from rich.markdown import Markdown as RichMarkdown
 
-@intercept_session_exceptions
+    RICH_AVAILABLE = True
+except ImportError:
+    RICH_AVAILABLE = False
+
+
+@session_class_wrapper
 class InspectAPI:
     def __init__(self, state: SessionState) -> None:
         self._state = state
@@ -25,7 +32,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."""
 
@@ -61,14 +68,19 @@ class InspectIssues:
             closest_match = set(difflib.get_close_matches(search, unique_types))
             issues = IssueList([issue for issue in issues if type(issue).__name__ in closest_match])
 
+        issue_str = "\n".join(
+            [f"  * **{type(issue).__name__}**: {issue.as_message(include_type=False)}" for issue in issues]
+        )
+        markdown_str = f"### {len(issues)} issues found\n\n{issue_str}"
+
         if IN_NOTEBOOK:
             from IPython.display import Markdown, display
 
-            issue_str = "\n".join(
-                [f"  * **{type(issue).__name__}**: {issue.as_message(include_type=False)}" for issue in issues]
-            )
-            message = f"### {len(issues)} issues found\n\n{issue_str}"
-            display(Markdown(message))
+            display(Markdown(markdown_str))
+        elif RICH_AVAILABLE:
+            from rich import print
+
+            print(RichMarkdown(markdown_str))
 
         if return_dataframe:
             return issues.to_pandas()
@@ -92,14 +104,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