chore(profiling): improve typing in profiler.py

Author: KowalskiThomas

ddtrace/internal/datadog/profiling/ddup/_ddup.pyi

@@ -1,4 +1,5 @@
 from typing import Dict
+from typing import Mapping
 from typing import Optional
 from typing import Union
 
@@ -12,7 +13,7 @@ def config(
     env: StringType = None,
     service: StringType = None,
     version: StringType = None,
-    tags: Optional[Dict[Union[str, bytes], Union[str, bytes]]] = None,
+    tags: Optional[Mapping[Union[str, bytes], Union[str, bytes]]] = None,
     max_nframes: Optional[int] = None,
     timeline_enabled: Optional[bool] = None,
     output_filename: Optional[str] = None,

ddtrace/internal/datadog/profiling/ddup/_ddup.pyx

@@ -2,7 +2,7 @@
 # cython: language_level=3
 
 import platform
-from typing import Dict
+from typing import Mapping
 from typing import Optional
 from typing import Union
 
@@ -327,7 +327,7 @@ def config(
         service: StringType = None,
         env: StringType = None,
         version: StringType = None,
-        tags: Optional[Dict[Union[str, bytes], Union[str, bytes]]] = None,
+        tags: Optional[Mapping[Union[str, bytes], Union[str, bytes]]] = None,
         max_nframes: Optional[int] = None,
         timeline_enabled: Optional[bool] = None,
         output_filename: StringType = None,

ddtrace/profiling/profiler.py

@@ -2,11 +2,14 @@
 import logging
 import os
 from typing import Any
+from typing import Callable
 from typing import Dict
-from typing import List  # noqa:F401
-from typing import Optional  # noqa:F401
-from typing import Type  # noqa:F401
-from typing import Union  # noqa:F401
+from typing import List
+from typing import Mapping
+from typing import Optional
+from typing import Type
+from typing import Union
+from typing import cast
 
 import ddtrace
 from ddtrace import config
@@ -29,8 +32,6 @@
 from ddtrace.profiling.collector import threading
 
 
-# TODO(vlad): add type annotations
-
 LOG = logging.getLogger(__name__)
 
 
@@ -42,10 +43,10 @@ class Profiler(object):
 
     """
 
-    def __init__(self, *args, **kwargs):
-        self._profiler = _ProfilerInstance(*args, **kwargs)
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        self._profiler: "_ProfilerInstance" = _ProfilerInstance(*args, **kwargs)
 
-    def start(self, stop_on_exit=True, profile_children=True):
+    def start(self, stop_on_exit: bool = True, profile_children: bool = True) -> None:
         """Start the profiler.
 
         :param stop_on_exit: Whether to stop the profiler and flush the profile on exit.
@@ -75,7 +76,7 @@ def start(self, stop_on_exit=True, profile_children=True):
 
         telemetry_writer.product_activated(TELEMETRY_APM_PRODUCT.PROFILER, True)
 
-    def stop(self, flush=True):
+    def stop(self, flush: bool = True) -> None:
         """Stop the profiler.
 
         :param flush: Flush last profile.
@@ -88,7 +89,7 @@ def stop(self, flush=True):
             # Not a best practice, but for backward API compatibility that allowed to call `stop` multiple times.
             pass
 
-    def _restart_on_fork(self):
+    def _restart_on_fork(self) -> None:
         # Be sure to stop the parent first, since it might have to e.g. unpatch functions
         # Do not flush data as we don't want to have multiple copies of the parent profile exported.
         try:
@@ -99,11 +100,7 @@ def _restart_on_fork(self):
         self._profiler = self._profiler.copy()
         self._profiler.start()
 
-    def __getattr__(
-        self,
-        key,  # type: str
-    ):
-        # type: (...) -> Any
+    def __getattr__(self, key: str) -> Any:
         return getattr(self._profiler, key)
 
 
@@ -145,22 +142,22 @@ def __init__(
         self.endpoint_collection_enabled: bool = endpoint_collection_enabled
 
         # Non-user-supplied values
-        self._collectors: List[Union[stack.StackCollector, memalloc.MemoryCollector]] = []
-        self._collectors_on_import: Any = None
+        self._collectors: List[collector.Collector] = []
+        self._collectors_on_import: Optional[List[tuple[str, Callable[[Any], None]]]] = None
         self._scheduler: Optional[Union[scheduler.Scheduler, scheduler.ServerlessScheduler]] = None
         self._lambda_function_name: Optional[str] = os.environ.get("AWS_LAMBDA_FUNCTION_NAME")
 
         self.__post_init__()
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         for k, v in vars(self).items():
             if k.startswith("_") or k in self._COPY_IGNORE_ATTRIBUTES:
                 continue
             if v != getattr(other, k, None):
                 return False
         return True
 
-    def _build_default_exporters(self):
+    def _build_default_exporters(self) -> None:
         if self._lambda_function_name is not None:
             self.tags.update({"functionname": self._lambda_function_name})
 
@@ -176,7 +173,7 @@ def _build_default_exporters(self):
             env=self.env,
             service=self.service,
             version=self.version,
-            tags=self.tags,
+            tags=cast(Mapping[Union[str, bytes], Union[str, bytes]], self.tags),
             max_nframes=profiling_config.max_frames,
             timeline_enabled=profiling_config.timeline_enabled,
             output_filename=profiling_config.output_pprof,
@@ -185,9 +182,7 @@ def _build_default_exporters(self):
         )
         ddup.start()
 
-    def __post_init__(self):
-        # type: (...) -> None
-
+    def __post_init__(self) -> None:
         if self._stack_collector_enabled:
             LOG.debug("Profiling collector (stack) enabled")
             try:
@@ -199,7 +194,7 @@ def __post_init__(self):
         if self._lock_collector_enabled:
             # These collectors require the import of modules, so we create them
             # if their import is detected at runtime.
-            def start_collector(collector_class: Type) -> None:
+            def start_collector(collector_class: Type[collector.Collector]) -> None:
                 with self._service_lock:
                     col = collector_class(tracer=self.tracer)
 
@@ -228,7 +223,7 @@ def start_collector(collector_class: Type) -> None:
 
         if self._pytorch_collector_enabled:
 
-            def start_collector(collector_class: Type) -> None:
+            def start_collector(collector_class: Type[collector.Collector]) -> None:
                 with self._service_lock:
                     col = collector_class()
 
@@ -254,18 +249,20 @@ def start_collector(collector_class: Type) -> None:
                 ModuleWatchdog.register_module_hook(module, hook)
 
         if self._memory_collector_enabled:
-            self._collectors.append(memalloc.MemoryCollector())
+            self._collectors.append(memalloc.MemoryCollector())  # type: ignore[arg-type]
 
         self._build_default_exporters()
 
-        scheduler_class = scheduler.ServerlessScheduler if self._lambda_function_name else scheduler.Scheduler  # type: (Type[Union[scheduler.Scheduler, scheduler.ServerlessScheduler]])
+        scheduler_class: Type[Union[scheduler.Scheduler, scheduler.ServerlessScheduler]] = (
+            scheduler.ServerlessScheduler if self._lambda_function_name else scheduler.Scheduler
+        )
 
         self._scheduler = scheduler_class(
             before_flush=self._collectors_snapshot,
             tracer=self.tracer,
         )
 
-    def _collectors_snapshot(self):
+    def _collectors_snapshot(self) -> None:
         for c in self._collectors:
             try:
                 c.snapshot()
@@ -274,7 +271,7 @@ def _collectors_snapshot(self):
 
     _COPY_IGNORE_ATTRIBUTES = {"status"}
 
-    def copy(self):
+    def copy(self) -> "_ProfilerInstance":
         return self.__class__(
             **{
                 key: value
@@ -283,8 +280,7 @@ def copy(self):
             }
         )
 
-    def _start_service(self):
-        # type: (...) -> None
+    def _start_service(self) -> None:
         """Start the profiler."""
         collectors = []
         for col in self._collectors:
@@ -301,14 +297,13 @@ def _start_service(self):
         if self._scheduler is not None:
             self._scheduler.start()
 
-    def _stop_service(self, flush=True, join=True):
-        # type: (bool, bool) -> None
+    def _stop_service(self, flush: bool = True, join: bool = True) -> None:
         """Stop the profiler.
 
         :param flush: Flush a last profile.
         """
         # Prevent doing more initialisation now that we are shutting down.
-        if self._lock_collector_enabled:
+        if self._lock_collector_enabled and self._collectors_on_import:
             for module, hook in self._collectors_on_import:
                 try:
                     ModuleWatchdog.unregister_module_hook(module, hook)