Add ANN rule for ruff

cognite/extractorutils/_inner_util.py

@@ -37,14 +37,14 @@ def resolve_log_level_for_httpx(level: str) -> str:
 
 
 class _DecimalEncoder(json.JSONEncoder):
-    def default(self, obj: Any) -> dict[str, str]:
+    def default(self, obj: Any) -> dict[str, str]:  # noqa: ANN401
         if isinstance(obj, Decimal):
             return {"type": "decimal_encoded", "value": str(obj)}
         return super().default(obj)
 
 
 class _DecimalDecoder(json.JSONDecoder):
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
         json.JSONDecoder.__init__(self, *args, object_hook=self.object_hook, **kwargs)
 
     def object_hook(self, obj_dict: dict[str, str]) -> dict[str, str] | Decimal:

cognite/extractorutils/base.py

@@ -112,7 +112,7 @@ def __init__(
         reload_config_interval: int | None = 300,
         reload_config_action: ReloadConfigAction = ReloadConfigAction.DO_NOTHING,
         success_message: str = "Successful shutdown",
-    ):
+    ) -> None:
         self.name = name
         self.description = description
         self.run_handle = run_handle

cognite/extractorutils/configtools/elements.py

@@ -29,6 +29,7 @@
 
 import yaml
 from prometheus_client import REGISTRY, start_http_server
+from typing_extensions import Self
 
 from cognite.client import ClientConfig, CogniteClient
 from cognite.client.credentials import (
@@ -926,7 +927,7 @@ class CastableInt(int):
     file.
     """
 
-    def __new__(cls, value: Any) -> "CastableInt":
+    def __new__(cls, value: int | str | bytes) -> Self:
         """
         Returns value as is if it's int.
 
@@ -955,7 +956,7 @@ class PortNumber(CastableInt):
     not a valid port number raises a ValueError at instantiation.
     """
 
-    def __new__(cls, value: Any) -> "PortNumber":
+    def __new__(cls, value: int | str | bytes) -> Self:
         """
         Try to cast the value to an integer and validate it as a port number.
 

cognite/extractorutils/configtools/loaders.py

@@ -72,7 +72,7 @@ class KeyVaultLoader:
         config: A dictionary containing the configuration for the keyvault.
     """
 
-    def __init__(self, config: dict | None):
+    def __init__(self, config: dict | None) -> None:
         self.config = config
 
         self.client: SecretClient | None = None
@@ -374,7 +374,7 @@ class ConfigResolver(Generic[CustomConfigClass]):
     Automatically reloads the configuration file if it has changed
     """
 
-    def __init__(self, config_path: str, config_type: type[CustomConfigClass]):
+    def __init__(self, config_path: str, config_type: type[CustomConfigClass]) -> None:
         self.config_path = config_path
         self.config_type = config_type
 

cognite/extractorutils/exceptions.py

@@ -25,7 +25,7 @@ class InvalidConfigError(Exception):
       * Unknown fields
     """
 
-    def __init__(self, message: str, details: list[str] | None = None):
+    def __init__(self, message: str, details: list[str] | None = None) -> None:
         super().__init__()
         self.message = message
         self.details = details

cognite/extractorutils/metrics.py

@@ -67,7 +67,7 @@ def __init__(self):
 T = TypeVar("T")
 
 
-def safe_get(cls: type[T], *args: Any, **kwargs: Any) -> T:
+def safe_get(cls: type[T], *args: Any, **kwargs: Any) -> T:  # noqa: ANN401
     """
     A factory for instances of metrics collections.
 
@@ -122,7 +122,7 @@ class BaseMetrics:
         process_scrape_interval: Interval (in seconds) between each fetch of data for the ``process_*`` gauges
     """
 
-    def __init__(self, extractor_name: str, extractor_version: str, process_scrape_interval: float = 15):
+    def __init__(self, extractor_name: str, extractor_version: str, process_scrape_interval: float = 15) -> None:
         extractor_name = extractor_name.strip().replace(" ", "_")
 
         self.startup = Gauge(f"{extractor_name}_start_time", "Timestamp (seconds) of when the extractor last started")
@@ -187,7 +187,7 @@ def __init__(
         push_interval: int | None = None,
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         self.push_interval = push_interval
         self.thread_name = thread_name
 
@@ -274,7 +274,7 @@ def __init__(
         password: str | None = None,
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         super().__init__(push_interval, thread_name, cancellation_token)
 
         self.username = username
@@ -283,7 +283,9 @@ def __init__(
 
         self.url = url
 
-    def _auth_handler(self, url: str, method: str, timeout: int, headers: list[tuple[str, str]], data: Any) -> Callable:
+    def _auth_handler(
+        self, url: str, method: str, timeout: int, headers: list[tuple[str, str]], data: bytes
+    ) -> Callable[[], None]:
         """
         Returns a authentication handler against the Prometheus Pushgateway to use in the pushadd_to_gateway method.
 
@@ -350,7 +352,7 @@ def __init__(
         data_set: EitherId | None = None,
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         super().__init__(push_interval, thread_name, cancellation_token)
 
         self.cdf_client = cdf_client

cognite/extractorutils/statestore/watermark.py

@@ -1,3 +1,6 @@
+# ruff: noqa: ANN401
+# TODO: the state stores should be generic over the type of state, not just Any.
+
 #  Copyright 2020 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -66,7 +69,7 @@ def __init__(
         trigger_log_level: str = "DEBUG",
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         super().__init__(
             save_interval=save_interval,
             trigger_log_level=trigger_log_level,
@@ -249,7 +252,7 @@ def __init__(
         trigger_log_level: str = "DEBUG",
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         super().__init__(save_interval, trigger_log_level, thread_name, cancellation_token)
 
         self._cdf_client = cdf_client
@@ -395,7 +398,7 @@ def __init__(
         trigger_log_level: str = "DEBUG",
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         super().__init__(save_interval, trigger_log_level, thread_name, cancellation_token)
 
         self._file_path = file_path

cognite/extractorutils/threading.py

@@ -6,7 +6,7 @@
 import signal
 from threading import Condition
 from time import time
-from typing import Any
+from types import FrameType
 
 
 class CancellationToken:
@@ -114,7 +114,7 @@ def cancel_on_interrupt(self) -> None:
         This will set the cancellation token instead of throwing a KeyboardInterrupt exception.
         """
 
-        def sigint_handler(sig_num: int, frame: Any) -> None:
+        def sigint_handler(sig_num: int, frame: FrameType | None) -> None:
             logger = logging.getLogger(__name__)
             logger.warning("Interrupt signal received, stopping extractor gracefully")
             self.cancel()

cognite/extractorutils/unstable/configuration/exceptions.py

@@ -12,7 +12,7 @@ class InvalidConfigError(Exception):
       * Unknown fields
     """
 
-    def __init__(self, message: str, details: list[str] | None = None):
+    def __init__(self, message: str, details: list[str] | None = None) -> None:
         super().__init__()
         self.message = message
         self.details = details

cognite/extractorutils/unstable/configuration/models.py

@@ -59,7 +59,7 @@ def __init__(self, scopes: str) -> None:
         self._scopes = list(scopes.split(" "))
 
     @classmethod
-    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema:
+    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema:  # noqa: ANN401
         return core_schema.no_info_after_validator_function(cls, handler(str))
 
     def __eq__(self, other: object) -> bool:
@@ -106,7 +106,7 @@ def __init__(self, expression: str) -> None:
         self._interval, self._expression = TimeIntervalConfig._parse_expression(expression)
 
     @classmethod
-    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema:
+    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema:  # noqa: ANN401
         """
         Pydantic hook to define how this class should be serialized/deserialized.
 

cognite/extractorutils/unstable/core/_dto.py

@@ -17,14 +17,14 @@ class CogniteModel(BaseModel):
       * exclude Nones from serialized JSON instead of having nulls in the response text.
     """
 
-    def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
+    def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]:  # noqa: ANN401
         if kwargs:
             kwargs["exclude_none"] = True
         else:
             kwargs = {"exclude_none": True}
         return BaseModel.model_dump(self, *args, **kwargs)
 
-    def dict(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
+    def dict(self, *args: Any, **kwargs: Any) -> dict[str, Any]:  # noqa: ANN401
         return self.model_dump(*args, **kwargs)
 
     model_config = ConfigDict(alias_generator=camelize, populate_by_name=True, extra="forbid")

cognite/extractorutils/unstable/core/tasks.py

@@ -28,7 +28,7 @@ class TaskContext(CogniteLogger):
     This class is used to log errors and messages related to the task execution.
     """
 
-    def __init__(self, task: "Task", extractor: "Extractor"):
+    def __init__(self, task: "Task", extractor: "Extractor") -> None:
         super().__init__()
         self._task = task
         self._extractor = extractor
@@ -87,7 +87,7 @@ def __init__(
         target: TaskTarget,
         description: str | None = None,
         schedule: ScheduleConfig,
-    ):
+    ) -> None:
         super().__init__(name=name, target=target, description=description)
         self.schedule = schedule
 

cognite/extractorutils/uploader/_base.py

@@ -50,7 +50,7 @@ def __init__(
         trigger_log_level: str = "DEBUG",
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         self.cdf_client = cdf_client
 
         self.threshold = max_queue_size if max_queue_size is not None else -1

cognite/extractorutils/uploader/assets.py

@@ -64,7 +64,7 @@ def __init__(
         trigger_log_level: str = "DEBUG",
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         super().__init__(
             cdf_client,
             post_upload_function,

cognite/extractorutils/uploader/data_modeling.py

@@ -50,7 +50,7 @@ def __init__(
         auto_create_start_nodes: bool = True,
         auto_create_end_nodes: bool = True,
         auto_create_direct_relations: bool = True,
-    ):
+    ) -> None:
         super().__init__(
             cdf_client,
             post_upload_function,

cognite/extractorutils/uploader/events.py

@@ -62,7 +62,7 @@ def __init__(
         trigger_log_level: str = "DEBUG",
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         # Super sets post_upload and threshold
         super().__init__(
             cdf_client,

cognite/extractorutils/uploader/files.py

@@ -103,13 +103,13 @@ def tell(self) -> int:
     # resolve the same way. These four useless methods with liberal use of Any are
     # 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:
+    def writelines(self, __lines: Any) -> None:  # noqa: ANN401
         """
         Not supported for ChunkedStream.
         """
         raise NotImplementedError()
 
-    def write(self, __b: Any) -> int:
+    def write(self, __b: Any) -> int:  # noqa: ANN401
         """
         Not supported for ChunkedStream.
         """
@@ -250,7 +250,7 @@ def __init__(
         max_parallelism: int | None = None,
         failure_logging_path: None | str = None,
         ssl_verify: bool | str = True,
-    ):
+    ) -> None:
         # Super sets post_upload and threshold
         super().__init__(
             cdf_client,
@@ -696,7 +696,7 @@ def __init__(
         overwrite_existing: bool = False,
         cancellation_token: CancellationToken | None = None,
         ssl_verify: bool | str = True,
-    ):
+    ) -> None:
         # Super sets post_upload and threshold
         super().__init__(
             cdf_client=cdf_client,

cognite/extractorutils/uploader/raw.py

@@ -67,7 +67,7 @@ def __init__(
         trigger_log_level: str = "DEBUG",
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         # Super sets post_upload and thresholds
         super().__init__(
             cdf_client,

cognite/extractorutils/uploader/time_series.py

@@ -122,7 +122,7 @@ def __init__(
         trigger_log_level: str = "DEBUG",
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         # Super sets post_upload and threshold
         super().__init__(
             cdf_client,
@@ -248,7 +248,7 @@ def __init__(
         create_missing: Callable[[str, DataPointList], TimeSeries] | bool = False,
         data_set_id: int | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         # Super sets post_upload and threshold
         super().__init__(
             cdf_client,
@@ -429,7 +429,7 @@ def __init__(
         create_missing: Callable[[NodeId, DataPointList], CogniteExtractorTimeSeriesApply] | bool = False,
         cancellation_token: CancellationToken | None = None,
         source: DirectRelationReference | None = None,
-    ):
+    ) -> None:
         super().__init__(
             cdf_client,
             post_upload_function,
@@ -636,7 +636,7 @@ def __init__(
         thread_name: str | None = None,
         create_missing: bool = False,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         # Super sets post_upload and threshold
         super().__init__(
             cdf_client,

cognite/extractorutils/uploader_extractor.py

@@ -105,7 +105,7 @@ def __init__(
         heartbeat_waiting_time: int = 600,
         handle_interrupts: bool = True,
         middleware: list[Callable[[dict], dict]] | None = None,
-    ):
+    ) -> None:
         super().__init__(
             name=name,
             description=description,
@@ -165,7 +165,7 @@ def handle_output(self, output: CdfTypes) -> None:
         else:
             raise ValueError(f"Unexpected type: {type(peek)}")
 
-    def _apply_middleware(self, item: Any) -> Any:
+    def _apply_middleware(self, item: Any) -> Any:  # noqa: ANN401
         for mw in self.middleware:
             item = mw(item)
         return item