chore(profiling): improve typing in _gevent.py

Author: KowalskiThomas

ddtrace/profiling/_gevent.py

@@ -13,18 +13,18 @@
 
 
 # Original objects
-_gevent_hub_spawn_raw = gevent.hub.spawn_raw
-_gevent_joinall = gevent.joinall
-_gevent_wait = gevent.wait
-_gevent_iwait = gevent.iwait
+_gevent_hub_spawn_raw: t.Callable[..., _Greenlet] = gevent.hub.spawn_raw
+_gevent_joinall: t.Callable[..., None] = gevent.joinall
+_gevent_wait: t.Callable[..., t.Any] = gevent.wait
+_gevent_iwait: t.Callable[..., t.Any] = gevent.iwait
 
 # Global package state
 _greenlet_frames: t.Dict[int, t.Union[FrameType, bool, None]] = {}
 _original_greenlet_tracer: t.Optional[t.Callable[[str, t.Any], None]] = None
 _greenlet_parent_map: t.Dict[int, int] = {}
 _parent_greenlet_count: t.Dict[int, int] = {}
 
-FRAME_NOT_SET = False  # Sentinel for when the frame is not set
+FRAME_NOT_SET: bool = False  # Sentinel for when the frame is not set
 
 
 class GreenletTrackingError(Exception):
@@ -34,7 +34,7 @@ class GreenletTrackingError(Exception):
 
 
 def track_gevent_greenlet(greenlet: _Greenlet) -> _Greenlet:
-    greenlet_id = thread.get_ident(greenlet)
+    greenlet_id: int = thread.get_ident(greenlet)
     frame: t.Union[FrameType, bool, None] = FRAME_NOT_SET
 
     try:
@@ -66,7 +66,7 @@ def update_greenlet_frame(greenlet_id: int, frame: t.Union[FrameType, bool, None
 def greenlet_tracer(event: str, args: t.Any) -> None:
     if event in {"switch", "throw"}:
         # This tracer function runs in the context of the target
-        origin, target = t.cast(t.Tuple[Greenlet, Greenlet], args)
+        origin, target = t.cast(t.Tuple[_Greenlet, _Greenlet], args)
 
         if (origin_id := thread.get_ident(origin)) not in _greenlet_frames:
             try:
@@ -105,7 +105,7 @@ def greenlet_tracer(event: str, args: t.Any) -> None:
 
 
 def untrack_greenlet(greenlet: _Greenlet) -> None:
-    greenlet_id = thread.get_ident(greenlet)
+    greenlet_id: int = thread.get_ident(greenlet)
     stack_v2.untrack_greenlet(greenlet_id)
     _greenlet_frames.pop(greenlet_id, None)
     _parent_greenlet_count.pop(greenlet_id, None)
@@ -140,8 +140,8 @@ def spawn_later(cls, *args: t.Any, **kwargs: t.Any) -> _Greenlet:
             return greenlet
 
     def join(self, *args: t.Any, **kwargs: t.Any) -> None:
-        target_id = thread.get_ident(self)
-        origin_id = thread.get_ident(gevent.getcurrent())
+        target_id: int = thread.get_ident(self)
+        origin_id: int = thread.get_ident(gevent.getcurrent())
 
         link_greenlets(target_id, origin_id)
 
@@ -165,7 +165,7 @@ def joinall(greenlets: t.Sequence[_Greenlet], *args: t.Any, **kwargs: t.Any) ->
     current_greenlet = gevent.getcurrent()
     if isinstance(current_greenlet, greenlet):
         current_greenlet = gevent.hub.get_hub()
-    current_greenlet_id = thread.get_ident(current_greenlet)
+    current_greenlet_id: int = thread.get_ident(current_greenlet)
     for g in greenlets:
         link_greenlets(thread.get_ident(g), current_greenlet_id)
     _gevent_joinall(greenlets, *args, **kwargs)
@@ -182,7 +182,7 @@ def _(*args: t.Any, **kwargs: t.Any) -> t.Any:
             current_greenlet = gevent.getcurrent()
             if isinstance(current_greenlet, greenlet):
                 current_greenlet = gevent.hub.get_hub()
-            current_greenlet_id = thread.get_ident(current_greenlet)
+            current_greenlet_id: int = thread.get_ident(current_greenlet)
             for g in greenlets:
                 link_greenlets(thread.get_ident(g), current_greenlet_id)
 
@@ -204,7 +204,7 @@ def patch() -> None:
 
     gevent.hub.spawn_raw = wrap_spawn(_gevent_hub_spawn_raw)
 
-    _original_greenlet_tracer = settrace(greenlet_tracer)
+    _original_greenlet_tracer = t.cast(t.Callable[[str, t.Any], None], settrace(greenlet_tracer))
 
 
 def unpatch() -> None: