PR review changes : Basemetrics subclass

Author: Hmnt39

cognite/extractorutils/metrics.py

@@ -458,3 +458,6 @@ def stop(self) -> None:
         self._push_to_server()
         self.upload_queue.stop()
         self.cancellation_token.cancel()
+
+
+MetricsType = TypeVar("MetricsType", bound=BaseMetrics)

cognite/extractorutils/unstable/core/base.py

@@ -69,7 +69,7 @@ def my_task_function(self, task_context: TaskContext) -> None:
 from typing_extensions import Self, assert_never
 
 from cognite.extractorutils._inner_util import _resolve_log_level
-from cognite.extractorutils.metrics import BaseMetrics, safe_get
+from cognite.extractorutils.metrics import BaseMetrics, MetricsType, safe_get
 from cognite.extractorutils.statestore import (
     AbstractStateStore,
     LocalStateStore,
@@ -127,13 +127,13 @@ def __init__(
         application_config: _T,
         current_config_revision: ConfigRevision,
         log_level_override: str | None = None,
-        metrics_class: type[BaseMetrics] | None = None,
+        metrics_class: type[MetricsType] | None = None,
     ) -> None:
         self.connection_config = connection_config
         self.application_config = application_config
         self.current_config_revision: ConfigRevision = current_config_revision
         self.log_level_override = log_level_override
-        self.metrics_class: type[BaseMetrics] | None = metrics_class
+        self.metrics_class: type[MetricsType] | None = metrics_class
 
 
 class Extractor(Generic[ConfigType], CogniteLogger):
@@ -273,17 +273,15 @@ def _setup_logging(self) -> None:
                             "Defaulted to console logging."
                         )
 
-    def _load_metrics(self, metrics_class: type[BaseMetrics] | None = None) -> BaseMetrics:
+    def _load_metrics(self, metrics_class: type[MetricsType] | None = None) -> MetricsType | BaseMetrics:
         """
         Loads metrics based on the provided metrics class.
 
         Reuses existing singleton if available to avoid Prometheus registry conflicts.
         """
         if metrics_class:
-            metrics_instance = safe_get(metrics_class)
-        else:
-            metrics_instance = safe_get(BaseMetrics, extractor_name=self.EXTERNAL_ID, extractor_version=self.VERSION)
-        return metrics_instance
+            return safe_get(metrics_class)
+        return safe_get(BaseMetrics, extractor_name=self.EXTERNAL_ID, extractor_version=self.VERSION)
 
     def _load_state_store(self) -> None:
         """

cognite/extractorutils/unstable/core/runtime.py

@@ -47,7 +47,7 @@ def main() -> None:
     CogniteAuthError,
     CogniteConnectionError,
 )
-from cognite.extractorutils.metrics import BaseMetrics
+from cognite.extractorutils.metrics import BaseMetrics, MetricsType
 from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.unstable.configuration.exceptions import InvalidArgumentError, InvalidConfigError
 from cognite.extractorutils.unstable.configuration.loaders import (
@@ -135,7 +135,7 @@ class Runtime(Generic[ExtractorType]):
     def __init__(
         self,
         extractor: type[ExtractorType],
-        metrics: type[BaseMetrics] | None = None,
+        metrics: type[MetricsType] | None = None,
     ) -> None:
         self._extractor_class = extractor
         self._cancellation_token = CancellationToken()