Merge pull request #635 from aiven/aiven-anton/chore/annotate-statsd

chore: Add type-hints to statsd modules

karapace/config.py

@@ -10,8 +10,8 @@
 from karapace.constants import DEFAULT_SCHEMA_TOPIC
 from karapace.utils import json_decode, json_encode, JSONDecodeError
 from pathlib import Path
-from typing import IO
-from typing_extensions import TypedDict
+from typing import IO, Mapping
+from typing_extensions import NotRequired, TypedDict
 
 import logging
 import os
@@ -76,6 +76,9 @@ class Config(TypedDict):
     master_election_strategy: str
     protobuf_runtime_directory: str
 
+    sentry: NotRequired[Mapping[str, object]]
+    tags: NotRequired[Mapping[str, object]]
+
 
 class ConfigDefaults(Config, total=False):
     ...

karapace/sentry/__init__.py

@@ -1,12 +1,13 @@
-from karapace.sentry.sentry_client_api import SentryClientAPI, SentryNoOpClient
-from typing import Dict, Optional
+from __future__ import annotations
+
+from karapace.sentry.sentry_client_api import KarapaceSentryConfig, SentryClientAPI, SentryNoOpClient
 
 import logging
 
 LOG = logging.getLogger(__name__)
 
 
-def _get_sentry_noop_client(sentry_config: Optional[Dict]) -> SentryClientAPI:
+def _get_sentry_noop_client(sentry_config: KarapaceSentryConfig) -> SentryClientAPI:
     return SentryNoOpClient(sentry_config=sentry_config)
 
 
@@ -17,13 +18,13 @@ def _get_sentry_noop_client(sentry_config: Optional[Dict]) -> SentryClientAPI:
     from karapace.sentry.sentry_client import SentryClient
 
     # If Sentry SDK can be imported in SentryClient the Sentry SDK can be initialized.
-    def _get_actual_sentry_client(sentry_config: Optional[Dict]) -> SentryClientAPI:
+    def _get_actual_sentry_client(sentry_config: KarapaceSentryConfig) -> SentryClientAPI:
         return SentryClient(sentry_config=sentry_config)
 
     _get_sentry_client = _get_actual_sentry_client
 except ImportError:
     LOG.warning("Cannot enable Sentry.io sending: importing 'sentry_sdk' failed")
 
 
-def get_sentry_client(sentry_config: Optional[Dict]) -> SentryClientAPI:
+def get_sentry_client(sentry_config: KarapaceSentryConfig) -> SentryClientAPI:
     return _get_sentry_client(sentry_config=sentry_config)

karapace/sentry/sentry_client.py

@@ -2,21 +2,23 @@
 Copyright (c) 2023 Aiven Ltd
 See LICENSE for details
 """
-from karapace.sentry.sentry_client_api import SentryClientAPI
-from typing import Dict, Optional
+from __future__ import annotations
+
+from karapace.sentry.sentry_client_api import KarapaceSentryConfig, SentryClientAPI
+from typing import Mapping
 
 # The Sentry SDK is optional, omit pylint import error
 import sentry_sdk
 
 
 class SentryClient(SentryClientAPI):
-    def __init__(self, sentry_config: Optional[Dict]) -> None:
+    def __init__(self, sentry_config: KarapaceSentryConfig) -> None:
         super().__init__(sentry_config=sentry_config)
         self._initialize_sentry()
 
     def _initialize_sentry(self) -> None:
         sentry_config = (
-            self.sentry_config.copy()
+            dict(self.sentry_config)
             if self.sentry_config is not None
             else {
                 "ignore_errors": [
@@ -34,7 +36,7 @@ def _initialize_sentry(self) -> None:
 
         # If the DSN is not in the config or in SENTRY_DSN environment variable
         # the Sentry client does not send any events.
-        sentry_sdk.init(**sentry_config)
+        sentry_sdk.init(**sentry_config)  # type: ignore[arg-type]
 
         # Don't send library logged errors to Sentry as there is also proper return value or raised exception to calling code
         from sentry_sdk.integrations.logging import ignore_logger
@@ -44,7 +46,12 @@ def _initialize_sentry(self) -> None:
         ignore_logger("kafka")
         ignore_logger("kafka.*")
 
-    def unexpected_exception(self, error: Exception, where: str, tags: Optional[Dict] = None) -> None:
+    def unexpected_exception(
+        self,
+        error: Exception,
+        where: str,
+        tags: Mapping[str, str] | None = None,
+    ) -> None:
         scope_args = {"tags": {"where": where, **(tags or {})}}
         sentry_sdk.Hub.current.capture_exception(error=error, scope=None, **scope_args)
 

karapace/sentry/sentry_client_api.py

@@ -2,14 +2,24 @@
 Copyright (c) 2023 Aiven Ltd
 See LICENSE for details
 """
-from typing import Dict, Optional
+from __future__ import annotations
+
+from typing import Mapping
+from typing_extensions import TypeAlias
+
+KarapaceSentryConfig: TypeAlias = "Mapping[str, object] | None"
 
 
 class SentryClientAPI:
-    def __init__(self, sentry_config: Optional[Dict]) -> None:
+    def __init__(self, sentry_config: KarapaceSentryConfig) -> None:
         self.sentry_config = sentry_config or {}
 
-    def unexpected_exception(self, error: Exception, where: str, tags: Optional[Dict] = None) -> None:
+    def unexpected_exception(
+        self,
+        error: Exception,
+        where: str,
+        tags: Mapping[str, str] | None = None,
+    ) -> None:
         pass
 
     def close(self) -> None:

karapace/statsd.py

@@ -13,29 +13,29 @@
 from contextlib import contextmanager
 from karapace.config import Config
 from karapace.sentry import get_sentry_client
-from typing import Any, Iterator
+from typing import Any, Final, Iterator
 
 import datetime
 import logging
 import socket
 import time
 
-STATSD_HOST = "127.0.0.1"
-STATSD_PORT = 8125
+STATSD_HOST: Final = "127.0.0.1"
+STATSD_PORT: Final = 8125
 LOG = logging.getLogger(__name__)
 
 
 class StatsClient:
     def __init__(
         self,
+        config: Config,
         host: str = STATSD_HOST,
         port: int = STATSD_PORT,
-        config: Config | None = None,
     ) -> None:
-        self._dest_addr = (host, port)
-        self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-        self._tags = config.get("tags", {})
-        self.sentry_client = get_sentry_client(sentry_config=config.get("sentry", None))
+        self._dest_addr: Final = (host, port)
+        self._socket: Final = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self._tags: Final = config.get("tags", {})
+        self.sentry_client: Final = get_sentry_client(sentry_config=config.get("sentry", None))
 
     @contextmanager
     def timing_manager(self, metric: str, tags: dict | None = None) -> Iterator[None]:
@@ -70,7 +70,7 @@ def _send(self, metric: str, metric_type: bytes, value: Any, tags: dict | None)
         try:
             # format: "user.logins,service=payroll,region=us-west:1|c"
             parts = [metric.encode("utf-8"), b":", str(value).encode("utf-8"), b"|", metric_type]
-            send_tags = self._tags.copy()
+            send_tags = dict(self._tags)
             send_tags.update(tags or {})
             for tag, tag_value in sorted(send_tags.items()):
                 if tag_value is None:

mypy.ini

@@ -85,9 +85,6 @@ ignore_errors = True
 [mypy-karapace.config]
 ignore_errors = True
 
-[mypy-karapace.statsd]
-ignore_errors = True
-
 [mypy-karapace.utils]
 ignore_errors = True