explicitly import all functions (#1413)

Co-authored-by: Håkon V. Treider <[email protected]>

cognite/client/_api/assets.py

@@ -29,7 +29,6 @@
 
 from typing_extensions import TypeAlias
 
-from cognite.client import utils
 from cognite.client._api_client import APIClient
 from cognite.client._constants import DEFAULT_LIMIT_READ
 from cognite.client.data_classes import (
@@ -49,8 +48,8 @@
 from cognite.client.data_classes.filters import Filter, _validate_filter
 from cognite.client.data_classes.shared import AggregateBucketResult
 from cognite.client.exceptions import CogniteAPIError
-from cognite.client.utils._auxiliary import split_into_chunks, split_into_n_parts
-from cognite.client.utils._concurrency import classify_error, get_priority_executor
+from cognite.client.utils._auxiliary import assert_type, split_into_chunks, split_into_n_parts
+from cognite.client.utils._concurrency import classify_error, execute_tasks, get_priority_executor
 from cognite.client.utils._identifier import IdentifierSequence
 from cognite.client.utils._text import to_camel_case
 from cognite.client.utils._validation import prepare_filter_sort, process_asset_subtree_ids, process_data_set_ids
@@ -578,7 +577,7 @@ def create(self, asset: Asset | Sequence[Asset]) -> Asset | AssetList:
                 >>> asset = Asset(name="my_pump", labels=[Label(external_id="PUMP")])
                 >>> res = c.assets.create(asset)
         """
-        utils._auxiliary.assert_type(asset, "asset", [Asset, Sequence])
+        assert_type(asset, "asset", [Asset, Sequence])
         return self._create_multiple(list_cls=AssetList, resource_cls=Asset, items=asset)
 
     def create_hierarchy(
@@ -708,7 +707,7 @@ def create_hierarchy(
             raise ValueError(f"'upsert_mode' must be either 'patch' or 'replace', not {upsert_mode!r}")
 
         if not isinstance(assets, AssetHierarchy):
-            utils._auxiliary.assert_type(assets, "assets", [Sequence])
+            assert_type(assets, "assets", [Sequence])
             assets = AssetHierarchy(assets, ignore_orphans=True)
 
         return _AssetHierarchyCreator(assets, assets_api=self).create(upsert, upsert_mode)
@@ -1016,7 +1015,7 @@ def _get_asset_subtree(self, assets: list, current_depth: int, depth: int | None
     def _get_children(self, assets: list) -> list:
         ids = [a.id for a in assets]
         tasks = [{"parent_ids": chunk, "limit": -1} for chunk in split_into_chunks(ids, 100)]
-        tasks_summary = utils._concurrency.execute_tasks(self.list, tasks=tasks, max_workers=self._config.max_workers)
+        tasks_summary = execute_tasks(self.list, tasks=tasks, max_workers=self._config.max_workers)
         tasks_summary.raise_compound_exception_if_failed_tasks()
         res_list = tasks_summary.results
         children = []

cognite/client/_api/data_modeling/instances.py

@@ -697,7 +697,7 @@ def search(
 
                 >>> from cognite.client import CogniteClient
                 >>> from cognite.client.data_classes.data_modeling import ViewId
-                >>> import cognite.client.data_classes.filters as filters
+                >>> from cognite.client.data_classes import filters
                 >>> c = CogniteClient()
                 >>> born_after_1970 = filters.Range(["mySpace", "PersonView/v1", "birthYear"], gt=1970)
                 >>> res = c.data_modeling.instances.search(ViewId("mySpace", "PersonView", "v1"),

cognite/client/_api/extractionpipelines.py

@@ -2,7 +2,6 @@
 
 from typing import TYPE_CHECKING, Any, Literal, Sequence, overload
 
-from cognite.client import utils
 from cognite.client._api_client import APIClient
 from cognite.client._constants import DEFAULT_LIMIT_READ
 from cognite.client.data_classes import (
@@ -17,6 +16,7 @@
     TimestampRange,
 )
 from cognite.client.data_classes.extractionpipelines import StringFilter
+from cognite.client.utils._auxiliary import assert_type
 from cognite.client.utils._identifier import IdentifierSequence
 
 if TYPE_CHECKING:
@@ -151,7 +151,7 @@ def create(
                 >>> extpipes = [ExtractionPipeline(name="extPipe1",...), ExtractionPipeline(name="extPipe2",...)]
                 >>> res = c.extraction_pipelines.create(extpipes)
         """
-        utils._auxiliary.assert_type(extraction_pipeline, "extraction_pipeline", [ExtractionPipeline, Sequence])
+        assert_type(extraction_pipeline, "extraction_pipeline", [ExtractionPipeline, Sequence])
         return self._create_multiple(
             list_cls=ExtractionPipelineList, resource_cls=ExtractionPipeline, items=extraction_pipeline
         )
@@ -303,7 +303,7 @@ def create(
                 >>> res = c.extraction_pipelines.runs.create(
                 ...     ExtractionPipelineRun(status="success", extpipe_external_id="extId"))
         """
-        utils._auxiliary.assert_type(run, "run", [ExtractionPipelineRun, Sequence])
+        assert_type(run, "run", [ExtractionPipelineRun, Sequence])
         return self._create_multiple(list_cls=ExtractionPipelineRunList, resource_cls=ExtractionPipelineRun, items=run)
 
 

cognite/client/_api/files.py

@@ -16,7 +16,6 @@
     overload,
 )
 
-from cognite.client import utils
 from cognite.client._api_client import APIClient
 from cognite.client._constants import _RUNNING_IN_BROWSER, DEFAULT_LIMIT_READ
 from cognite.client.data_classes import (
@@ -33,6 +32,7 @@
 )
 from cognite.client.exceptions import CogniteFileUploadError
 from cognite.client.utils._auxiliary import find_duplicates
+from cognite.client.utils._concurrency import execute_tasks
 from cognite.client.utils._identifier import Identifier, IdentifierSequence
 from cognite.client.utils._validation import process_asset_subtree_ids, process_data_set_ids
 
@@ -480,9 +480,7 @@ def upload(
                         file_metadata = copy.copy(file_metadata)
                         file_metadata.name = file_name
                         tasks.append((file_metadata, file_path, overwrite))
-            tasks_summary = utils._concurrency.execute_tasks(
-                self._upload_file_from_path, tasks, self._config.max_workers
-            )
+            tasks_summary = execute_tasks(self._upload_file_from_path, tasks, self._config.max_workers)
             tasks_summary.raise_compound_exception_if_failed_tasks(task_unwrap_fn=lambda x: x[0].name)
             return FileMetadataList(tasks_summary.results)
         raise ValueError(f"The path '{path}' does not exist")
@@ -607,7 +605,7 @@ def retrieve_download_urls(
             dict(url_path="/files/downloadlink", json={"items": id_batch}, params=query_params)
             for id_batch in id_batches
         ]
-        tasks_summary = utils._concurrency.execute_tasks(self._post, tasks, max_workers=self._config.max_workers)
+        tasks_summary = execute_tasks(self._post, tasks, max_workers=self._config.max_workers)
         tasks_summary.raise_compound_exception_if_failed_tasks()
         results = tasks_summary.joined_results(unwrap_fn=lambda res: res.json()["items"])
         return {result.get("id") or result["externalId"]: result["downloadUrl"] for result in results}
@@ -766,9 +764,7 @@ def _download_files_to_directory(
     ) -> None:
         self._warn_on_duplicate_filenames(filepaths)
         tasks = [(directory, id, id_to_metadata, filepath) for id, filepath in zip(all_ids, filepaths)]
-        tasks_summary = utils._concurrency.execute_tasks(
-            self._process_file_download, tasks, max_workers=self._config.max_workers
-        )
+        tasks_summary = execute_tasks(self._process_file_download, tasks, max_workers=self._config.max_workers)
         tasks_summary.raise_compound_exception_if_failed_tasks(
             task_unwrap_fn=lambda task: id_to_metadata[IdentifierSequence.unwrap_identifier(task[1])],
             str_format_element_fn=lambda metadata: metadata.id,

cognite/client/_api/functions.py

@@ -15,7 +15,6 @@
 from typing import TYPE_CHECKING, Any, Callable, Sequence, cast
 from zipfile import ZipFile
 
-from cognite.client import utils
 from cognite.client._api_client import APIClient
 from cognite.client._constants import DEFAULT_LIMIT_READ
 from cognite.client.data_classes import (
@@ -33,7 +32,7 @@
     TimestampRange,
 )
 from cognite.client.data_classes.functions import FunctionCallsFilter, FunctionsStatus
-from cognite.client.utils._auxiliary import is_unlimited
+from cognite.client.utils._auxiliary import assert_type, is_unlimited, local_import
 from cognite.client.utils._identifier import Identifier, IdentifierSequence
 from cognite.client.utils._session import create_session_and_return_nonce
 
@@ -187,8 +186,8 @@ def create(
         elif function_handle:
             _validate_function_handle(function_handle)
             file_id = self._zip_and_upload_handle(function_handle, name, external_id)
-        utils._auxiliary.assert_type(cpu, "cpu", [Number], allow_none=True)
-        utils._auxiliary.assert_type(memory, "memory", [Number], allow_none=True)
+        assert_type(cpu, "cpu", [Number], allow_none=True)
+        assert_type(memory, "memory", [Number], allow_none=True)
 
         sleep_time = 1.0  # seconds
         for i in range(MAX_RETRIES):
@@ -350,8 +349,8 @@ def retrieve_multiple(
                 >>> c = CogniteClient()
                 >>> res = c.functions.retrieve_multiple(external_ids=["func1", "func2"])
         """
-        utils._auxiliary.assert_type(ids, "id", [Sequence], allow_none=True)
-        utils._auxiliary.assert_type(external_ids, "external_id", [Sequence], allow_none=True)
+        assert_type(ids, "id", [Sequence], allow_none=True)
+        assert_type(external_ids, "external_id", [Sequence], allow_none=True)
         return self._retrieve_multiple(
             identifiers=IdentifierSequence.load(ids=ids, external_ids=external_ids),
             resource_cls=Function,
@@ -665,7 +664,7 @@ def _validate_and_parse_requirements(requirements: list[str]) -> list[str]:
     Returns:
         list[str]: The parsed requirements
     """
-    constructors = cast(Any, utils._auxiliary.local_import("pip._internal.req.constructors"))
+    constructors = cast(Any, local_import("pip._internal.req.constructors"))
     install_req_from_line = constructors.install_req_from_line
     parsed_reqs: list[str] = []
     for req in requirements: