chore(profiling): improve typing in memalloc.py

Author: KowalskiThomas

ddtrace/profiling/collector/memalloc.py

@@ -9,6 +9,8 @@
 from typing import Type
 from typing import cast
 
+from typing_extensions import Self
+
 
 try:
     from ddtrace.profiling.collector import _memalloc
@@ -55,8 +57,9 @@ def start(self) -> None:
             _memalloc.stop()
             _memalloc.start(self.max_nframe, self.heap_sample_size)
 
-    def __enter__(self) -> None:
+    def __enter__(self) -> Self:
         self.start()
+        return self
 
     def __exit__(
         self,
@@ -66,7 +69,7 @@ def __exit__(
     ) -> None:
         self.stop()
 
-    def join(self) -> None:
+    def join(self, timeout: Optional[float] = None) -> None:
         pass
 
     def stop(self) -> None:
@@ -76,7 +79,8 @@ def stop(self) -> None:
             except RuntimeError:
                 LOG.debug("Failed to stop memalloc profiling on shutdown", exc_info=True)
 
-    def _get_thread_id_ignore_set(self) -> Set[int]:
+    @staticmethod
+    def _get_thread_id_ignore_set() -> Set[int]:
         # This method is not perfect and prone to race condition in theory, but very little in practice.
         # Anyhow it's not a big deal — it's a best effort feature.
         return {
@@ -86,7 +90,8 @@ def _get_thread_id_ignore_set(self) -> Set[int]:
         }
 
     def snapshot(self) -> None:
-        thread_id_ignore_set = self._get_thread_id_ignore_set()
+        """Take a snapshot of collected data, to be exported."""
+        thread_id_ignore_set = MemoryCollector._get_thread_id_ignore_set()
 
         try:
             if _memalloc is None:

ddtrace/profiling/profiler.py

@@ -142,7 +142,9 @@ def __init__(
         self.endpoint_collection_enabled: bool = endpoint_collection_enabled
 
         # Non-user-supplied values
-        self._collectors: List[collector.Collector] = []
+        # Note: memalloc.MemoryCollector is not a subclass of collector.Collector, so we need to use a union type.
+        #       This is because its snapshot method cannot be static.
+        self._collectors: List[collector.Collector | memalloc.MemoryCollector] = []
         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")
@@ -251,7 +253,7 @@ def start_collector(collector_class: Type[collector.Collector]) -> None:
                 ModuleWatchdog.register_module_hook(module, hook)
 
         if self._memory_collector_enabled:
-            self._collectors.append(memalloc.MemoryCollector())  # type: ignore[arg-type]
+            self._collectors.append(memalloc.MemoryCollector())
 
         self._build_default_exporters()