TMP

Author: haakonvt

cognite/client/_basic_api_client.py

@@ -199,14 +199,14 @@ def _init_async_http_clients(self) -> None:
             refresh_auth_header=self._refresh_auth_header,
         )
 
-    def __getstate__(self):
+    def __getstate__(self) -> dict[str, Any]:
         """Prepare object for pickling by removing unpicklable async clients."""
         state = self.__dict__.copy()
         state.pop("_http_client")
         state.pop("_http_client_with_retry")
         return state
 
-    def __setstate__(self, state):
+    def __setstate__(self, state: dict[str, Any]) -> None:
         """Restore object after unpickling."""
         self.__dict__.update(state)
         self._init_async_http_clients()

cognite/client/_constants.py

@@ -2,6 +2,16 @@
 
 from typing import Literal
 
+__all__ = [
+    "DATA_MODELING_DEFAULT_LIMIT_READ",
+    "DEFAULT_DATAPOINTS_CHUNK_SIZE",
+    "DEFAULT_LIMIT_READ",
+    "MAX_VALID_INTERNAL_ID",
+    "NUMPY_IS_AVAILABLE",
+    "OMITTED",
+    "_RUNNING_IN_BROWSER",
+]
+
 try:
     from pyodide.ffi import IN_BROWSER as _RUNNING_IN_BROWSER  # type: ignore [import-not-found]
 except ModuleNotFoundError:

cognite/client/_sync_api/postgres_gateway/tables.py

@@ -120,7 +120,9 @@ def retrieve(
         """
         return run_sync(
             self.__async_client.postgres_gateway.tables.retrieve(
-                username=username, tablename=tablename, ignore_unknown_ids=ignore_unknown_ids
+                username=username,
+                tablename=tablename,
+                ignore_unknown_ids=ignore_unknown_ids,  # type: ignore[arg-type]
             )
         )
 

cognite/client/data_classes/datapoints.py

@@ -1411,3 +1411,21 @@ def to_pandas(  # type: ignore [override]
             include_status=include_status,
             include_unit=include_unit,
         )
+
+
+__all__ = [
+    "ALL_SORTED_DP_AGGS",
+    "Datapoint",
+    "DatapointList",
+    "DatapointsQuery",
+    "DatapointsQueryResult",
+    "DatapointsQueryResultList",
+    "MaxDatapoint",
+    "MaxDatapointWithStatus",
+    "MinDatapoint",
+    "MinDatapointWithStatus",
+    "NumericDatapoint",
+    "NumericDatapointList",
+    "StringDatapoint",
+    "StringDatapointList",
+]

cognite/client/data_classes/files.py

@@ -613,6 +613,7 @@ async def __aexit__(
         self._check_errors_before_completing(exc_type)
         await self._cognite_client.files._complete_multipart_upload(self)
         self._upload_is_finalized = True
+        return None
 
     def __enter__(self) -> Self:
         if self._upload_is_finalized:
@@ -621,9 +622,9 @@ def __enter__(self) -> Self:
 
     def __exit__(
         self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
-    ) -> bool | None:
+    ) -> None:
         if exc_type is not None:
-            return False
+            return
 
         self._check_errors_before_completing(exc_type)
         run_sync(self._cognite_client.files._complete_multipart_upload(self))

cognite/client/data_classes/workflows.py

@@ -1389,7 +1389,7 @@ def __init__(
     ) -> None:
         super().__init__(with_, select)
         # Parameters and cursors are not supported for workflow trigger queries:
-        self.parameters = None
+        self.parameters = None  # type: ignore [assignment]
         self.cursors = None  # type: ignore [assignment]
 
     @classmethod

cognite/client/utils/_pandas_helpers.py

@@ -25,6 +25,8 @@
 from cognite.client.utils._time import TIME_ATTRIBUTES
 
 if TYPE_CHECKING:
+    import numpy as np
+    import numpy.typing as npt
     import pandas as pd
 
     from cognite.client.data_classes import Datapoints, DatapointsArray, DatapointsArrayList, DatapointsList
@@ -200,7 +202,7 @@ def convert_dps_to_dataframe(
     include_status: bool,
     include_unit: bool,
     include_errors: bool = False,  # old leftover misuse of Datapoints class :(
-):
+) -> pd.DataFrame:
     pd = local_import("pandas")
     columns = _extract_column_info_from_dps_for_dataframe(
         dps, include_status=include_status, include_errors=include_errors
@@ -262,7 +264,7 @@ def as_array(self) -> NumpyObjArray | NumpyFloat64Array | NumpyInt64Array | Nump
         else:
             return self._convert_to_array_for_agg_dps()
 
-    def _convert_to_array_for_raw_dps(self):
+    def _convert_to_array_for_raw_dps(self) -> npt.NDArray[np.object_]:
         import numpy as np
 
         match self.is_string, self.status_info:
@@ -279,7 +281,7 @@ def _convert_to_array_for_raw_dps(self):
                 # the datapoints object themselves:
                 assert_never(f"Invalid combination of is_string={self.is_string} and status_info={self.status_info}")
 
-    def _convert_to_array_for_agg_dps(self):
+    def _convert_to_array_for_agg_dps(self) -> npt.NDArray[np.object_]:
         import numpy as np
 
         from cognite.client.utils._datapoints import ensure_int_numpy

cognite/client/utils/_time.py

@@ -346,3 +346,20 @@ def split_time_range(start: int, end: int, n_splits: int, granularity_in_ms: int
     # Find a `delta_ms` that's a multiple of granularity in ms (trivial for raw queries).
     delta_ms = granularity_in_ms * round(tot_ms / n_splits / granularity_in_ms)
     return [*(start + delta_ms * i for i in range(n_splits)), end]
+
+
+__all__ = [
+    "MAX_TIMESTAMP_MS",
+    "MIN_TIMESTAMP_MS",
+    "ZoneInfo",
+    "convert_and_isoformat_time_attrs",
+    "convert_data_modeling_timestamp",
+    "datetime_to_ms",
+    "datetime_to_ms_iso_timestamp",
+    "granularity_to_ms",
+    "ms_to_datetime",
+    "parse_str_timezone",
+    "parse_str_timezone_offset",
+    "split_time_range",
+    "timestamp_to_ms",
+]