chore(profiling): improve typing in profiler.py

Author: KowalskiThomas

ddtrace/profiling/profiler.py

@@ -2,11 +2,12 @@
 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 Optional
+from typing import Type
+from typing import Union
 
 import ddtrace
 from ddtrace import config
@@ -28,9 +29,6 @@
 from ddtrace.profiling.collector import stack
 from ddtrace.profiling.collector import threading
 
-
-# TODO(vlad): add type annotations
-
 LOG = logging.getLogger(__name__)
 
 
@@ -42,10 +40,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 +73,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 +86,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 +97,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 +139,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})
 
@@ -185,8 +179,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")
@@ -228,7 +221,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()
 
@@ -258,14 +251,16 @@ def start_collector(collector_class: Type) -> None:
 
         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 +269,7 @@ def _collectors_snapshot(self):
 
     _COPY_IGNORE_ATTRIBUTES = {"status"}
 
-    def copy(self):
+    def copy(self) -> "_ProfilerInstance":
         return self.__class__(
             **{
                 key: value
@@ -283,8 +278,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,8 +295,7 @@ 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.