add docstrings with url, args & examples

Author: haakonvt

cognite/client/_api/data_modeling/statistics.py

@@ -1,12 +1,16 @@
 from __future__ import annotations
 
 import itertools
-from collections.abc import Sequence
-from typing import TYPE_CHECKING
+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, ProjectStatsAndLimits
+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
@@ -21,23 +25,77 @@ def __init__(self, config: ClientConfig, api_version: str | None, cognite_client
         self._RETRIEVE_LIMIT = 100  # may need to be renamed, but fine for now
 
     def project(self) -> ProjectStatsAndLimits:
-        """Usage data and limits for a project's data modelling usage, including data model schemas and graph instances"""
+        """`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
         )
 
-    def per_space(self, space: str | Sequence[str] | None = None, return_all: bool = False) -> InstanceStatsList:
-        """Statistics and limits for all spaces in the project, or for a select subset"""
+    @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 not None:
-            # Statistics and limits for spaces by their space identifiers
-            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)
-                )
+        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)
             )
-        raise ValueError("Either 'space' or 'return_all' must be specified")
+        )

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)