chore: add missing typing for bytecode utilities

ddtrace/internal/bytecode_injection/core.py

@@ -199,7 +199,7 @@ def _inject_invocation_nonrecursive(
     instrumented_lines: list[int] = []
     is_first_instrumented_module_line = code.co_name == "<module>"
 
-    def append_instruction(opcode: int, extended_arg: int):
+    def append_instruction(opcode: int, extended_arg: int) -> None:
         """
         Append an operation and its argument to the new bytecode.
 

ddtrace/internal/wrapping/asyncs.py

@@ -1,4 +1,5 @@
 import sys
+from types import CodeType
 
 import bytecode as bc
 
@@ -710,7 +711,7 @@
     raise RuntimeError(msg)
 
 
-def wrap_async(instrs, code, lineno):
+def wrap_async(instrs: list[bc.Instr], code: CodeType, lineno: int) -> None:
     if (bc.CompilerFlags.ASYNC_GENERATOR | bc.CompilerFlags.COROUTINE) & code.co_flags:
         if ASYNC_HEAD_ASSEMBLY is not None:
             instrs[0:0] = ASYNC_HEAD_ASSEMBLY.bind()

ddtrace/internal/wrapping/context.py

@@ -391,7 +391,7 @@ class LazyWrappedFunction(Protocol):
 
     __dd_lazy_contexts__: list[WrappingContext]
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
         pass
 
 
@@ -421,11 +421,11 @@ def wrap(self) -> None:
                 super().wrap()
                 return
 
-            def trampoline(_, args, kwargs):
+            def trampoline(_: t.Any, args: tuple, kwargs: dict) -> t.Any:
                 with tl:
                     f = t.cast(WrappedFunction, self.__wrapped__)
                     if is_wrapped_with(self.__wrapped__, trampoline):
-                        f = unwrap(f, trampoline)
+                        f = t.cast(WrappedFunction, unwrap(f, trampoline))
 
                         self._trampoline = None
 
@@ -469,7 +469,7 @@ class ContextWrappedFunction(Protocol):
 
     __dd_context_wrapped__ = None  # type: t.Optional[_UniversalWrappingContext]
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
         pass
 
 
@@ -521,14 +521,19 @@ def __enter__(self) -> "_UniversalWrappingContext":
     def _exit(self) -> None:
         self.__exit__(*sys.exc_info())
 
-    def __exit__(self, *exc) -> None:
-        if exc == (None, None, None):
+    def __exit__(
+        self,
+        exc_type: t.Optional[type[BaseException]],
+        exc_value: t.Optional[BaseException],
+        traceback: t.Optional[TracebackType],
+    ) -> None:
+        if exc_value is None:
             return
 
         for context in self._contexts[::-1]:
-            context.__exit__(*exc)
+            context.__exit__(exc_type, exc_value, traceback)
 
-        super().__exit__(*exc)
+        super().__exit__(exc_type, exc_value, traceback)
 
     def __return__(self, value: T) -> T:
         for context in self._contexts[::-1]:

ddtrace/internal/wrapping/generators.py

@@ -1,4 +1,7 @@
 import sys
+from types import CodeType
+
+import bytecode as bc
 
 from ddtrace.internal.assembly import Assembly
 
@@ -481,7 +484,7 @@
     raise RuntimeError(msg)
 
 
-def wrap_generator(instrs, code, lineno):
+def wrap_generator(instrs: list[bc.Instr], code: CodeType, lineno: int) -> None:
     if GENERATOR_HEAD_ASSEMBLY is not None:
         instrs[0:0] = GENERATOR_HEAD_ASSEMBLY.bind(lineno=lineno)
 

mypy.ini

@@ -11,6 +11,14 @@ ignore_missing_imports = true
 namespace_packages = true
 plugins = envier.mypy
 
+[mypy-ddtrace.internal.wrapping.*]
+disallow_untyped_defs = true
+disallow_incomplete_defs = true
+
+[mypy-ddtrace.internal.bytecode_injection.*]
+disallow_untyped_defs = true
+disallow_incomplete_defs = true
+
 [mypy-ddtrace.debugging.*]
 disallow_untyped_defs = true
 disallow_incomplete_defs = true