Add /models/statistics

cognite/client/_api/data_modeling/__init__.py

@@ -7,6 +7,7 @@
 from cognite.client._api.data_modeling.graphql import DataModelingGraphQLAPI
 from cognite.client._api.data_modeling.instances import InstancesAPI
 from cognite.client._api.data_modeling.spaces import SpacesAPI
+from cognite.client._api.data_modeling.statistics import StatisticsAPI
 from cognite.client._api.data_modeling.views import ViewsAPI
 from cognite.client._api_client import APIClient
 
@@ -24,3 +25,4 @@ def __init__(self, config: ClientConfig, api_version: str | None, cognite_client
         self.views = ViewsAPI(config, api_version, cognite_client)
         self.instances = InstancesAPI(config, api_version, cognite_client)
         self.graphql = DataModelingGraphQLAPI(config, api_version, cognite_client)
+        self.statistics = StatisticsAPI(config, api_version, cognite_client)

cognite/client/_api/data_modeling/statistics.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import itertools
+from typing import TYPE_CHECKING, Any, Literal, overload
+
+from cognite.client._api_client import APIClient
+from cognite.client.data_classes.data_modeling.ids import _load_space_identifier
+from cognite.client.data_classes.data_modeling.statistics import (
+    InstanceStatsList,
+    InstanceStatsPerSpace,
+    ProjectStatsAndLimits,
+)
+from cognite.client.utils.useful_types import SequenceNotStr
+
+if TYPE_CHECKING:
+    from cognite.client._cognite_client import CogniteClient
+    from cognite.client.config import ClientConfig
+
+
+class StatisticsAPI(APIClient):
+    _RESOURCE_PATH = "/models/statistics"
+
+    def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None:
+        super().__init__(config, api_version, cognite_client)
+        self._RETRIEVE_LIMIT = 100  # may need to be renamed, but fine for now
+
+    def project(self) -> ProjectStatsAndLimits:
+        """`Retrieve project-wide usage data and limits <https://developer.cognite.com/api#tag/Statistics/operation/getStatistics>`_
+
+        Returns the usage data and limits for a project's data modelling usage, including data model schemas and graph instances
+
+        Returns:
+            ProjectStatsAndLimits: The requested statistics and limits
+
+        Examples:
+
+            Fetch project statistics (and limits) and check the current number of data models vs.
+            and how many more can be created:
+
+                >>> from cognite.client import CogniteClient
+                >>> client = CogniteClient()
+                >>> stats = client.data_modeling.statistics.project()
+                >>> num_dm = stats.data_models.current
+                >>> num_dm_left = stats.data_models.limit - num_dm
+        """
+        return ProjectStatsAndLimits._load(
+            self._get(self._RESOURCE_PATH).json(), project=self._cognite_client._config.project
+        )
+
+    @overload
+    def per_space(self, space: str, return_all: Literal[False]) -> InstanceStatsPerSpace: ...
+
+    @overload
+    def per_space(self, space: Any, return_all: Literal[True]) -> InstanceStatsList: ...
+
+    @overload
+    def per_space(self, space: SequenceNotStr[str], return_all: bool) -> InstanceStatsList: ...
+
+    def per_space(
+        self, space: str | SequenceNotStr[str] | None = None, return_all: bool = False
+    ) -> InstanceStatsPerSpace | InstanceStatsList:
+        """`Retrieve usage data and limits per space <https://developer.cognite.com/api#tag/Statistics/operation/getSpaceStatisticsByIds>`_
+
+        See also: `Retrieve statistics and limits for all spaces <https://developer.cognite.com/api#tag/Statistics/operation/getSpaceStatistics>`_
+
+        Args:
+            space (str | SequenceNotStr[str] | None): The space or spaces to retrieve statistics for.
+            return_all (bool): If True, fetch statistics for all spaces. If False, fetch statistics for the specified space(s).
+
+        Returns:
+            InstanceStatsPerSpace | InstanceStatsList: InstanceStatsPerSpace if a single space is given, else InstanceStatsList (which is a list of InstanceStatsPerSpace)
+
+        Examples:
+
+            Fetch statistics for a single space:
+
+                >>> from cognite.client import CogniteClient
+                >>> client = CogniteClient()
+                >>> res = client.data_modeling.statistics.per_space("my-space")
+
+            Fetch statistics for multiple spaces:
+                >>> res = client.data_modeling.statistics.per_space(
+                ...     ["my-space1", "my-space2"]
+                ... )
+
+            Fetch statistics for all spaces (ignores the 'space' argument):
+                >>> res = client.data_modeling.statistics.per_space(return_all=True)
+        """
+        if return_all:
+            return InstanceStatsList._load(self._get(self._RESOURCE_PATH + "/spaces").json()["items"])
+
+        elif space is None:
+            raise ValueError("Either 'space' or 'return_all' must be specified")
+
+        ids = _load_space_identifier(space)
+        return InstanceStatsList._load(
+            itertools.chain.from_iterable(
+                self._post(self._RESOURCE_PATH + "/spaces/byids", json={"items": chunk.as_dicts()}).json()["items"]
+                for chunk in ids.chunked(self._RETRIEVE_LIMIT)
+            )
+        )

cognite/client/data_classes/data_modeling/ids.py

@@ -190,7 +190,7 @@ def version(self) -> str | None: ...
 Id = tuple[str, str] | tuple[str, str, str] | IdLike | VersionedIdLike
 
 
-def _load_space_identifier(ids: str | SequenceNotStr[str]) -> DataModelingIdentifierSequence:
+def _load_space_identifier(ids: str | Sequence[str] | SequenceNotStr[str]) -> DataModelingIdentifierSequence:
     is_sequence = isinstance(ids, Sequence) and not isinstance(ids, str)
     spaces = [ids] if isinstance(ids, str) else ids
     return DataModelingIdentifierSequence(

cognite/client/data_classes/data_modeling/statistics.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+from collections import UserList
+from collections.abc import Iterable
+from dataclasses import asdict, dataclass
+from typing import TYPE_CHECKING, Any
+
+from typing_extensions import Self
+
+from cognite.client.utils._importing import local_import
+from cognite.client.utils._pandas_helpers import notebook_display_with_fallback
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+
+@dataclass
+class InstanceStatsPerSpace:
+    space: str
+    nodes: int
+    edges: int
+    soft_deleted_nodes: int
+    soft_deleted_edges: int
+
+    @classmethod
+    def _load(cls, data: dict[str, Any]) -> Self:
+        return cls(
+            space=data["space"],
+            nodes=data["nodes"],
+            edges=data["edges"],
+            soft_deleted_nodes=data["softDeletedNodes"],
+            soft_deleted_edges=data["softDeletedEdges"],
+        )
+
+    def _repr_html_(self) -> str:
+        return notebook_display_with_fallback(self)
+
+    def to_pandas(self) -> pd.DataFrame:
+        pd = local_import("pandas")
+        space = (dumped := asdict(self)).pop("space")
+        return pd.Series(dumped).to_frame(name=space)
+
+
+class InstanceStatsList(UserList):
+    def __init__(self, items: list[InstanceStatsPerSpace]):
+        super().__init__(items)
+
+    @classmethod
+    def _load(cls, data: Iterable[dict[str, Any]]) -> Self:
+        return cls([InstanceStatsPerSpace._load(item) for item in data])
+
+    def _repr_html_(self) -> str:
+        return notebook_display_with_fallback(self)
+
+    def to_pandas(self) -> pd.DataFrame:
+        pd = local_import("pandas")
+        df = pd.DataFrame([asdict(item) for item in self]).set_index("space")
+        order_by_total = (df["nodes"] + df["edges"]).sort_values(ascending=False).index
+        return df.loc[order_by_total]
+
+
+@dataclass
+class CountLimit:
+    count: int
+    limit: int
+
+    @classmethod
+    def _load(cls, data: dict[str, Any]) -> Self:
+        return cls(count=data["count"], limit=data["limit"])
+
+
+@dataclass
+class InstanceStatsAndLimits:
+    nodes: int
+    edges: int
+    instances: int
+    instances_limit: int
+    soft_deleted_nodes: int
+    soft_deleted_edges: int
+    soft_deleted_instances: int
+    soft_deleted_instances_limit: int
+
+    @classmethod
+    def _load(cls, data: dict[str, Any]) -> Self:
+        return cls(
+            nodes=data["nodes"],
+            edges=data["edges"],
+            instances=data["instances"],
+            instances_limit=data["instancesLimit"],
+            soft_deleted_nodes=data["softDeletedNodes"],
+            soft_deleted_edges=data["softDeletedEdges"],
+            soft_deleted_instances=data["softDeletedInstances"],
+            soft_deleted_instances_limit=data["softDeletedInstancesLimit"],
+        )
+
+    def _repr_html_(self) -> str:
+        return notebook_display_with_fallback(self)
+
+    def to_pandas(self) -> pd.DataFrame:
+        pd = local_import("pandas")
+        return pd.Series(asdict(self)).to_frame()
+
+
+@dataclass
+class ProjectStatsAndLimits:
+    project: str
+    spaces: CountLimit
+    containers: CountLimit
+    views: CountLimit
+    data_models: CountLimit
+    container_properties: CountLimit
+    instances: InstanceStatsAndLimits
+    concurrent_read_limit: int
+    concurrent_write_limit: int
+    concurrent_delete_limit: int
+
+    @classmethod
+    def _load(cls, data: dict[str, Any], project: str) -> Self:
+        return cls(
+            project=project,
+            spaces=CountLimit._load(data["spaces"]),
+            containers=CountLimit._load(data["containers"]),
+            views=CountLimit._load(data["views"]),
+            data_models=CountLimit._load(data["dataModels"]),
+            container_properties=CountLimit._load(data["containerProperties"]),
+            instances=InstanceStatsAndLimits._load(data["instances"]),
+            concurrent_read_limit=data["concurrentReadLimit"],
+            concurrent_write_limit=data["concurrentWriteLimit"],
+            concurrent_delete_limit=data["concurrentDeleteLimit"],
+        )
+
+    def _repr_html_(self) -> str:
+        return notebook_display_with_fallback(self)
+
+    def to_pandas(self) -> pd.DataFrame:
+        pd = local_import("pandas")
+        project = (dumped := asdict(self)).pop("project")
+        return pd.Series(dumped).to_frame(name=project)

cognite/client/testing.py

@@ -17,6 +17,7 @@
 from cognite.client._api.data_modeling.graphql import DataModelingGraphQLAPI
 from cognite.client._api.data_modeling.instances import InstancesAPI
 from cognite.client._api.data_modeling.spaces import SpacesAPI
+from cognite.client._api.data_modeling.statistics import StatisticsAPI
 from cognite.client._api.data_modeling.views import ViewsAPI
 from cognite.client._api.data_sets import DataSetsAPI
 from cognite.client._api.datapoints import DatapointsAPI
@@ -114,6 +115,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.data_modeling.views = MagicMock(spec_set=ViewsAPI)
         self.data_modeling.instances = MagicMock(spec_set=InstancesAPI)
         self.data_modeling.graphql = MagicMock(spec_set=DataModelingGraphQLAPI)
+        self.data_modeling.statistics = MagicMock(spec_set=StatisticsAPI)
 
         self.data_sets = MagicMock(spec_set=DataSetsAPI)
 

cognite/client/utils/_pandas_helpers.py

@@ -7,7 +7,7 @@
 from inspect import signature
 from itertools import chain
 from numbers import Integral
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, Protocol
 
 from cognite.client.exceptions import CogniteImportError
 from cognite.client.utils._importing import local_import
@@ -18,7 +18,6 @@
     import pandas as pd
 
     from cognite.client.data_classes import Datapoints, DatapointsArray, DatapointsArrayList, DatapointsList
-    from cognite.client.data_classes._base import T_CogniteResource, T_CogniteResourceList
 
 
 NULLABLE_INT_COLS = {
@@ -88,7 +87,11 @@ def concat_dps_dataframe_list(
     return concat_dataframes_with_nullable_int_cols(dfs)
 
 
-def notebook_display_with_fallback(inst: T_CogniteResource | T_CogniteResourceList, **kwargs: Any) -> str:
+class PandasConvertible(Protocol):
+    def to_pandas(self) -> pd.DataFrame: ...
+
+
+def notebook_display_with_fallback(inst: PandasConvertible, **kwargs: Any) -> str:
     params = signature(inst.to_pandas).parameters
     # Default of False enforced (when accepted by method):
     if "camel_case" in params:

tests/tests_unit/test_docstring_examples.py

@@ -25,7 +25,7 @@
     units,
     workflows,
 )
-from cognite.client._api.data_modeling import containers, data_models, graphql, instances, spaces, views
+from cognite.client._api.data_modeling import containers, data_models, graphql, instances, spaces, statistics, views
 from cognite.client._api.hosted_extractors import destinations, jobs, mappings, sources
 from cognite.client._api.postgres_gateway import tables as postgres_gateway_tables
 from cognite.client._api.postgres_gateway import users as postgres_gateway_users
@@ -114,6 +114,7 @@ def test_data_modeling(self):
         run_docstring_tests(data_models)
         run_docstring_tests(spaces)
         run_docstring_tests(graphql)
+        run_docstring_tests(statistics)
 
     def test_datapoint_subscriptions(self):
         run_docstring_tests(datapoints_subscriptions)