cognite-extractor-utils 7.5.13__py3-none-any.whl → 7.6.0__py3-none-any.whl

cognite/extractorutils/uploader/files.py

@@ -1,3 +1,7 @@
+"""
+Upload queue for files.
+"""
+
 #  Copyright 2023 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -88,6 +92,9 @@ class ChunkedStream(RawIOBase, BinaryIO):
         self._current_chunk_size = -1
 
     def tell(self) -> int:
+        """
+        Get the current position of the stream.
+        """
         return self._pos
 
     # RawIOBase is (stupidly) incompatible with BinaryIO
@@ -97,12 +104,21 @@ class ChunkedStream(RawIOBase, BinaryIO):
     # required to satisfy mypy.
     # This may be solvable by changing the typing in the python SDK to use typing.Protocol.
     def writelines(self, __lines: Any) -> None:
+        """
+        Not supported for ChunkedStream.
+        """
         raise NotImplementedError()
 
     def write(self, __b: Any) -> int:
+        """
+        Not supported for ChunkedStream.
+        """
         raise NotImplementedError()
 
     def __enter__(self) -> "ChunkedStream":
+        """
+        Wraps around the inner stream's ``__enter__`` method, for use as context manager.
+        """
         return super().__enter__()
 
     def __exit__(
@@ -111,14 +127,23 @@ class ChunkedStream(RawIOBase, BinaryIO):
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None:
+        """
+        Wraps around the inner stream's ``__exit__`` method, for use as context manager.
+        """
         return super().__exit__(exc_type, exc_val, exc_tb)
 
     @property
     def chunk_count(self) -> int:
+        """
+        Number of chunks in the stream.
+        """
         return ceil(self._stream_length / self._max_chunk_size)
 
     @property
     def len(self) -> int:
+        """
+        Length of the current chunk, in bytes.
+        """
         return len(self)
 
     @property
@@ -129,12 +154,21 @@ class ChunkedStream(RawIOBase, BinaryIO):
         return self._chunk_index
 
     def __len__(self) -> int:
+        """
+        Length of the current chunk, in bytes.
+        """
         return self._current_chunk_size
 
     def readable(self) -> bool:
+        """
+        Check if the stream is readable. Always True for ChunkedStream.
+        """
         return True
 
     def read(self, size: int = -1) -> bytes:
+        """
+        Read bytes from the current chunk.
+        """
         if size < 0:
             size = self._current_chunk_size - self._pos
 
@@ -162,6 +196,10 @@ class ChunkedStream(RawIOBase, BinaryIO):
 
 
 class IOByteStream(SyncByteStream):
+    """
+    Wraps a BinaryIO stream to be used as a httpx SyncByteStream.
+    """
+
     CHUNK_SIZE = 65_536
 
     def __init__(self, stream: BinaryIO) -> None:
@@ -169,6 +207,9 @@ class IOByteStream(SyncByteStream):
         self._is_stream_consumed = False
 
     def __iter__(self) -> Iterator[bytes]:
+        """
+        Iterate over the stream, yielding chunks of data.
+        """
         if self._is_stream_consumed:
             raise StreamConsumed()
         chunk = self._stream.read(self.CHUNK_SIZE)
@@ -179,7 +220,7 @@ class IOByteStream(SyncByteStream):
 
 class IOFileUploadQueue(AbstractUploadQueue):
     """
-    Upload queue for files using BinaryIO
+    Upload queue for files using BinaryIO.
 
     Note that if the upload fails, the stream needs to be restarted, so
     the enqueued callback needs to produce a new IO object for each call.
@@ -264,6 +305,9 @@ class IOFileUploadQueue(AbstractUploadQueue):
             _QUEUES += 1
 
     def initialize_failure_logging(self) -> None:
+        """
+        Initialize the failure logging manager if a path is provided in the constructor.
+        """
         self._file_failure_manager: FileFailureManager | None = (
             FileFailureManager(path_to_file=self.failure_logging_path)
             if self.failure_logging_path is not None
@@ -271,14 +315,23 @@ class IOFileUploadQueue(AbstractUploadQueue):
         )
 
     def get_failure_logger(self) -> FileFailureManager | None:
+        """
+        Get the failure logger for this upload queue, if it exists.
+        """
         return self._file_failure_manager
 
     def add_entry_failure_logger(self, file_name: str, error: Exception) -> None:
+        """
+        Add an entry to the failure logger if it exists.
+        """
         if self._file_failure_manager is not None:
             error_reason = str(error)
             self._file_failure_manager.add(file_name=file_name, error_reason=error_reason)
 
     def flush_failure_logger(self) -> None:
+        """
+        Flush the failure logger if it exists, writing all failures to the file.
+        """
         if self._file_failure_manager is not None:
             self.logger.info("Flushing failure logs")
             self._file_failure_manager.write_to_file()
@@ -313,14 +366,14 @@ class IOFileUploadQueue(AbstractUploadQueue):
             )
             if any_unchaged:
                 update = FileMetadataUpdate(external_id=file_meta.external_id)
-                any = False
+                need_update = False
                 if file_meta.source:
-                    any = True
+                    need_update = True
                     update.source.set(file_meta.source)
                 if file_meta.directory:
-                    any = True
+                    need_update = True
                     update.directory.set(file_meta.directory)
-                if any:
+                if need_update:
                     self.cdf_client.files.update(update)
 
         return file_meta_response, url
@@ -373,7 +426,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
     def _upload_multipart(self, size: int, file: BinaryIO, file_meta: FileMetadataOrCogniteExtractorFile) -> None:
         chunks = ChunkedStream(file, self.max_file_chunk_size, size)
         self.logger.debug(
-            f"File {file_meta.external_id} is larger than 5GiB ({size})" f", uploading in {chunks.chunk_count} chunks"
+            f"File {file_meta.external_id} is larger than 5GiB ({size}), uploading in {chunks.chunk_count} chunks"
         )
 
         returned_file_metadata = self._create_multi_part(file_meta, chunks)
@@ -428,13 +481,14 @@ class IOFileUploadQueue(AbstractUploadQueue):
         extra_retries: tuple[type[Exception], ...] | dict[type[Exception], Callable[[Any], bool]] | None = None,
     ) -> None:
         """
-        Add file to upload queue. The file will start uploading immedeately. If the size of the queue is larger than
-        the specified max size, this call will block until it's completed the upload.
+        Add file to upload queue.
+
+        The file will start uploading immediately. If the size of the queue is larger than the specified max size, this
+        call will block until it's completed the upload.
 
         Args:
             file_meta: File metadata-object
-            file_name: Path to file to be uploaded.
-                If none, the file object will still be created, but no data is uploaded
+            read_file: Callable that returns a BinaryIO stream to read the file from.
             extra_retries: Exception types that might be raised by ``read_file`` that should be retried
         """
         retries = cognite_exceptions()
@@ -568,7 +622,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
     def upload(self, fail_on_errors: bool = True, timeout: float | None = None) -> None:
         """
-        Wait for all uploads to finish
+        Wait for all uploads to finish.
         """
         for future in self.upload_queue:
             future.result(timeout=timeout)
@@ -581,7 +635,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
     def __enter__(self) -> "IOFileUploadQueue":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -598,7 +652,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
         exc_tb: TracebackType | None,
     ) -> None:
         """
-        Wraps around stop method, for use as context manager
+        Wraps around stop method, for use as context manager.
 
         Args:
             exc_type: Exception type
@@ -610,7 +664,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
     def __len__(self) -> int:
         """
-        The size of the upload queue
+        The size of the upload queue.
 
         Returns:
             Number of events in queue
@@ -620,7 +674,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
 class FileUploadQueue(IOFileUploadQueue):
     """
-    Upload queue for files
+    Upload queue for files.
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -661,8 +715,10 @@ class FileUploadQueue(IOFileUploadQueue):
         file_name: str | PathLike,
     ) -> None:
         """
-        Add file to upload queue. The queue will be uploaded if the queue size is larger than the threshold
-        specified in the __init__.
+        Add file to upload queue.
+
+        The file will start uploading immediately. If the size of the queue is larger than the specified max size, this
+        call will block until it's completed the upload.
 
         Args:
             file_meta: File metadata-object
@@ -678,7 +734,7 @@ class FileUploadQueue(IOFileUploadQueue):
 
 class BytesUploadQueue(IOFileUploadQueue):
     """
-    Upload queue for bytes
+    Upload queue for bytes.
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -714,11 +770,14 @@ class BytesUploadQueue(IOFileUploadQueue):
 
     def add_to_upload_queue(self, content: bytes, file_meta: FileMetadataOrCogniteExtractorFile) -> None:
         """
-        Add object to upload queue. The queue will be uploaded if the queue size is larger than the threshold
-        specified in the __init__.
+        Add file to upload queue.
+
+        The file will start uploading immediately. If the size of the queue is larger than the specified max size, this
+        call will block until it's completed the upload.
+
         Args:
             content: bytes object to upload
-            metadata: metadata for the given bytes object
+            file_meta: File metadata-object
         """
 
         def get_byte_io() -> BinaryIO:

cognite/extractorutils/uploader/raw.py

@@ -1,3 +1,7 @@
+"""
+Upload queue for RAW.
+"""
+
 #  Copyright 2023 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -41,7 +45,7 @@ from cognite.extractorutils.util import cognite_exceptions, retry
 
 class RawUploadQueue(AbstractUploadQueue):
     """
-    Upload queue for RAW
+    Upload queue for RAW.
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -84,8 +88,9 @@ class RawUploadQueue(AbstractUploadQueue):
 
     def add_to_upload_queue(self, database: str, table: str, raw_row: Row) -> None:
         """
-        Adds a row to the upload queue. The queue will be uploaded if the queue size is larger than the threshold
-        specified in the __init__.
+        Adds a row to the upload queue.
+
+        The queue will be uploaded if the queue size is larger than the threshold specified in the ``__init__``.
 
         Args:
             database: The database to upload the Raw object to
@@ -109,7 +114,7 @@ class RawUploadQueue(AbstractUploadQueue):
 
     def upload(self) -> None:
         """
-        Trigger an upload of the queue, clears queue afterwards
+        Trigger an upload of the queue, clears queue afterwards.
         """
 
         @retry(
@@ -154,7 +159,7 @@ class RawUploadQueue(AbstractUploadQueue):
 
     def __enter__(self) -> "RawUploadQueue":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -166,7 +171,7 @@ class RawUploadQueue(AbstractUploadQueue):
         self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None:
         """
-        Wraps around stop method, for use as context manager
+        Wraps around stop method, for use as context manager.
 
         Args:
             exc_type: Exception type
@@ -177,7 +182,7 @@ class RawUploadQueue(AbstractUploadQueue):
 
     def __len__(self) -> int:
         """
-        The size of the upload queue
+        The size of the upload queue.
 
         Returns:
             Number of elements in queue

cognite/extractorutils/uploader/time_series.py

@@ -1,3 +1,6 @@
+"""
+Upload queue for time series and sequences.
+"""
 #  Copyright 2023 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -81,7 +84,7 @@ def default_time_series_factory(external_id: str, datapoints: DataPointList) ->
 
 class TimeSeriesUploadQueue(AbstractUploadQueue):
     """
-    Upload queue for time series
+    Upload queue for time series.
 
     Args:
         cdf_client: Cognite Data Fusion client to use
@@ -141,7 +144,7 @@ class TimeSeriesUploadQueue(AbstractUploadQueue):
         self.data_set_id = data_set_id
 
     def _verify_datapoint_time(self, time: int | float | datetime | str) -> bool:
-        if isinstance(time, int) or isinstance(time, float):
+        if isinstance(time, int | float):
             return not math.isnan(time) and time >= MIN_DATAPOINT_TIMESTAMP
         elif isinstance(time, str):
             return False
@@ -155,10 +158,7 @@ class TimeSeriesUploadQueue(AbstractUploadQueue):
             )
         elif isinstance(value, str):
             return len(value) <= MAX_DATAPOINT_STRING_LENGTH
-        elif isinstance(value, datetime):
-            return False
-        else:
-            return True
+        return not isinstance(value, datetime)
 
     def _is_datapoint_valid(
         self,
@@ -172,11 +172,16 @@ class TimeSeriesUploadQueue(AbstractUploadQueue):
             return True
 
     def add_to_upload_queue(
-        self, *, id: int | None = None, external_id: str | None = None, datapoints: DataPointList | None = None
+        self,
+        *,
+        id: int | None = None,  # noqa: A002
+        external_id: str | None = None,
+        datapoints: DataPointList | None = None,
     ) -> None:
         """
-        Add data points to upload queue. The queue will be uploaded if the queue size is larger than the threshold
-        specified in the __init__.
+        Add data points to upload queue.
+
+        The queue will be uploaded if the queue size is larger than the threshold specified in the ``__init__``.
 
         Args:
             id: Internal ID of time series. Either this or external_id must be set.
@@ -209,7 +214,7 @@ class TimeSeriesUploadQueue(AbstractUploadQueue):
 
     def upload(self) -> None:
         """
-        Trigger an upload of the queue, clears queue afterwards
+        Trigger an upload of the queue, clears queue afterwards.
         """
 
         @retry(
@@ -239,9 +244,7 @@ class TimeSeriesUploadQueue(AbstractUploadQueue):
 
                 if self.create_missing:
                     # Get the time series that can be created
-                    create_these_ids = set(
-                        [id_dict["externalId"] for id_dict in ex.not_found if "externalId" in id_dict]
-                    )
+                    create_these_ids = {id_dict["externalId"] for id_dict in ex.not_found if "externalId" in id_dict}
                     datapoints_lists: dict[str, DataPointList] = {
                         ts_dict["externalId"]: ts_dict["datapoints"]
                         for ts_dict in upload_this
@@ -294,7 +297,7 @@ class TimeSeriesUploadQueue(AbstractUploadQueue):
                 ]
             )
 
-            for _either_id, datapoints in self.upload_queue.items():
+            for datapoints in self.upload_queue.values():
                 self.points_written.inc(len(datapoints))
 
             try:
@@ -309,7 +312,7 @@ class TimeSeriesUploadQueue(AbstractUploadQueue):
 
     def __enter__(self) -> "TimeSeriesUploadQueue":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -321,7 +324,7 @@ class TimeSeriesUploadQueue(AbstractUploadQueue):
         self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None:
         """
-        Wraps around stop method, for use as context manager
+        Wraps around stop method, for use as context manager.
 
         Args:
             exc_type: Exception type
@@ -332,7 +335,7 @@ class TimeSeriesUploadQueue(AbstractUploadQueue):
 
     def __len__(self) -> int:
         """
-        The size of the upload queue
+        The size of the upload queue.
 
         Returns:
             Number of data points in queue
@@ -341,6 +344,21 @@ class TimeSeriesUploadQueue(AbstractUploadQueue):
 
 
 class SequenceUploadQueue(AbstractUploadQueue):
+    """
+    Upload queue for sequences.
+
+    Args:
+        cdf_client: Cognite Data Fusion client to use
+        post_upload_function: A function that will be called after each upload. The function will be given one
+            argument: A list of the events that were uploaded.
+        max_queue_size: Maximum size of upload queue. Defaults to no max size.
+        max_upload_interval: Automatically trigger an upload each m seconds when run as a thread (use start/stop
+            methods).
+        trigger_log_level: Log level to log upload triggers to.
+        thread_name: Thread name of uploader thread.
+        create_missing: Create missing sequences if possible (ie, if external id is used).
+    """
+
     def __init__(
         self,
         cdf_client: CogniteClient,
@@ -352,19 +370,6 @@ class SequenceUploadQueue(AbstractUploadQueue):
         create_missing: bool = False,
         cancellation_token: CancellationToken | None = None,
     ):
-        """
-        Args:
-            cdf_client: Cognite Data Fusion client to use
-            post_upload_function: A function that will be called after each upload. The function will be given one
-                argument: A list of the events that were uploaded.
-            max_queue_size: Maximum size of upload queue. Defaults to no max size.
-            max_upload_interval: Automatically trigger an upload each m seconds when run as a thread (use start/stop
-                methods).
-            trigger_log_level: Log level to log upload triggers to.
-            thread_name: Thread name of uploader thread.
-            create_missing: Create missing sequences if possible (ie, if external id is used)
-        """
-
         # Super sets post_upload and threshold
         super().__init__(
             cdf_client,
@@ -393,7 +398,7 @@ class SequenceUploadQueue(AbstractUploadQueue):
     def set_sequence_metadata(
         self,
         metadata: dict[str, str | int | float],
-        id: int | None = None,
+        id: int | None = None,  # noqa: A002
         external_id: str | None = None,
         asset_external_id: str | None = None,
         dataset_external_id: str | None = None,
@@ -401,8 +406,10 @@ class SequenceUploadQueue(AbstractUploadQueue):
         description: str | None = None,
     ) -> None:
         """
-        Set sequence metadata. Metadata will be cached until the sequence is created. The metadata will be updated
-        if the sequence already exists
+        Set sequence metadata.
+
+        Metadata will be cached until the sequence is created. The metadata will be updated if the sequence already
+        exists.
 
         Args:
             metadata: Sequence metadata
@@ -427,10 +434,13 @@ class SequenceUploadQueue(AbstractUploadQueue):
             self.sequence_descriptions[either_id] = description
 
     def set_sequence_column_definition(
-        self, col_def: list[dict[str, str]], id: int | None = None, external_id: str | None = None
+        self,
+        col_def: list[dict[str, str]],
+        id: int | None = None,  # noqa: A002
+        external_id: str | None = None,
     ) -> None:
         """
-        Set sequence column definition
+        Set sequence column definition.
 
         Args:
             col_def: Sequence column definition
@@ -450,12 +460,13 @@ class SequenceUploadQueue(AbstractUploadQueue):
         | SequenceData
         | SequenceRows,
         column_external_ids: list[dict] | None = None,
-        id: int | None = None,
+        id: int | None = None,  # noqa: A002
         external_id: str | None = None,
     ) -> None:
         """
-        Add sequence rows to upload queue. Mirrors implementation of SequenceApi.insert. Inserted rows will be
-        cached until uploaded
+        Add sequence rows to upload queue.
+
+        Mirrors implementation of SequenceApi.insert. Inserted rows will be cached until uploaded.
 
         Args:
             rows: The rows to be inserted. Can either be a list of tuples, a list of ["rownumber": ..., "values": ...]
@@ -466,7 +477,6 @@ class SequenceUploadQueue(AbstractUploadQueue):
             external_id: Sequence external ID
                 Us if id is None
         """
-
         if len(rows) == 0:
             pass
 
@@ -509,7 +519,7 @@ class SequenceUploadQueue(AbstractUploadQueue):
 
     def upload(self) -> None:
         """
-        Trigger an upload of the queue, clears queue afterwards
+        Trigger an upload of the queue, clears queue afterwards.
         """
 
         @retry(
@@ -571,15 +581,14 @@ class SequenceUploadQueue(AbstractUploadQueue):
 
     def _create_or_update(self, either_id: EitherId) -> None:
         """
-        Create or update sequence, based on provided metadata and column definitions
+        Create or update sequence, based on provided metadata and column definitions.
 
         Args:
             either_id: Id/External Id of sequence to be updated
         """
-
         column_def = self.column_definitions.get(either_id)
         if column_def is None:
-            self.logger.error(f"Can't create sequence {str(either_id)}, no column definitions provided")
+            self.logger.error(f"Can't create sequence {either_id!s}, no column definitions provided")
 
         try:
             seq = self.cdf_client.sequences.create(
@@ -596,7 +605,7 @@ class SequenceUploadQueue(AbstractUploadQueue):
             )
 
         except CogniteDuplicatedError:
-            self.logger.info(f"Sequnce already exist: {either_id}")
+            self.logger.info(f"Sequence already exist: {either_id}")
             seq = self.cdf_client.sequences.retrieve(  # type: ignore [assignment]
                 id=either_id.internal_id,
                 external_id=either_id.external_id,
@@ -608,7 +617,7 @@ class SequenceUploadQueue(AbstractUploadQueue):
 
     def _resolve_asset_ids(self) -> None:
         """
-        Resolve id of assets if specified, for use in sequence creation
+        Resolve id of assets if specified, for use in sequence creation.
         """
         assets = set(self.sequence_asset_external_ids.values())
         assets.discard(None)  # type: ignore  # safeguard, remove Nones if any
@@ -628,7 +637,7 @@ class SequenceUploadQueue(AbstractUploadQueue):
 
     def _resolve_dataset_ids(self) -> None:
         """
-        Resolve id of datasets if specified, for use in sequence creation
+        Resolve id of datasets if specified, for use in sequence creation.
         """
         datasets = set(self.sequence_dataset_external_ids.values())
         datasets.discard(None)  # type: ignore  # safeguard, remove Nones if any
@@ -648,7 +657,7 @@ class SequenceUploadQueue(AbstractUploadQueue):
 
     def __enter__(self) -> "SequenceUploadQueue":
         """
-        Wraps around start method, for use as context manager
+        Wraps around start method, for use as context manager.
 
         Returns:
             self
@@ -660,7 +669,7 @@ class SequenceUploadQueue(AbstractUploadQueue):
         self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None:
         """
-        Wraps around stop method, for use as context manager
+        Wraps around stop method, for use as context manager.
 
         Args:
             exc_type: Exception type
@@ -671,7 +680,7 @@ class SequenceUploadQueue(AbstractUploadQueue):
 
     def __len__(self) -> int:
         """
-        The size of the upload queue
+        The size of the upload queue.
 
         Returns:
             Number of data points in queue