chore: bytecode strict typing

ddtrace/debugging/_expressions.py

@@ -122,7 +122,7 @@ def _make_function(self, instrs: list[Instr], args: tuple[str, ...], name: str)
         abstract_code = Bytecode([*instrs, Instr("RETURN_VALUE")])
 
         abstract_code.argcount = len(args)
-        abstract_code.argnames = args
+        abstract_code.argnames = list(args)
         abstract_code.name = name
 
         if sys.version_info >= (3, 11):
@@ -188,7 +188,7 @@ def _compile_arg_predicate(self, ast: DDASTType) -> Optional[list[Instr]]:
                 raise ValueError("Invalid argument: %r" % b)
 
             short_circuit = Label()
-            return ca + short_circuit_instrs(_type, short_circuit) + cb + [short_circuit]
+            return ca + short_circuit_instrs(_type, short_circuit) + cb + cast(list[Instr], [short_circuit])
 
         if _type in {"eq", "ge", "gt", "le", "lt", "ne"}:
             a, b = args

ddtrace/internal/assembly.py

@@ -68,7 +68,7 @@ def __init__(
     ) -> None:
         self._labels: dict[str, bc.Label] = {}
         self._ref_labels: dict[str, bc.Label] = {}
-        self._tb: t.Optional[bc.TryBegin] = None
+        self._tb: t.Optional[bc.TryBegin] = None  # type: ignore[name-defined]
         self._instrs = bc.Bytecode()
         self._instrs.name = name or "<assembly>"
         self._instrs.filename = filename or __file__
@@ -95,7 +95,8 @@ def parse_label(self, line: str) -> t.Optional[bc.Label]:
         if label_ident in self._labels:
             raise ValueError("label %s already defined" % label_ident)
 
-        label = self._labels[label_ident] = self._ref_labels.pop(label_ident, None) or bc.Label()
+        existing = self._ref_labels.pop(label_ident, None)
+        label = self._labels[label_ident] = existing if existing is not None else bc.Label()
 
         return label
 
@@ -177,7 +178,7 @@ def parse_expr(self, text: str) -> t.Any:
 
     def parse_opcode_arg(self, text: str) -> t.Union[bc.Label, str, int, t.Any]:
         if not text:
-            return bc.UNSET
+            return bc.UNSET  # type: ignore[attr-defined]
 
         return (
             self.parse_label_ref(text)
@@ -272,10 +273,10 @@ def dis(self) -> None:
                 print(f"    {entry.name:<32}{{{entry.arg}}}")
             elif isinstance(entry, bc.Label):
                 print(f"{self._label_ident(entry)}:")
-            elif isinstance(entry, bc.TryBegin):
+            elif isinstance(entry, bc.TryBegin):  # type: ignore[attr-defined]
                 print(f"try @{self._label_ident(entry.target)} (lasti={entry.push_lasti})")
 
-    def __iter__(self) -> t.Iterator[bc.Instr]:
+    def __iter__(self) -> t.Iterator[t.Any]:
         return iter(self._instrs)
 
     def __len__(self) -> int:

ddtrace/internal/bytecode_injection/__init__.py

@@ -4,6 +4,7 @@
 from typing import Callable  # noqa:F401
 
 from bytecode import Bytecode
+from bytecode import Instr
 
 from ddtrace.internal.assembly import Assembly
 from ddtrace.internal.compat import PYTHON_VERSION_INFO as PY
@@ -91,19 +92,18 @@ def _inject_hook(code: Bytecode, hook: HookType, lineno: int, arg: Any) -> None:
     last_lineno = None
     instrs = set()
     for i, instr in enumerate(code):
-        try:
-            if instr.lineno == last_lineno:
-                continue
-            last_lineno = instr.lineno
-            # Some lines might be implemented across multiple instruction
-            # offsets, and sometimes a NOP is used as a placeholder. We skip
-            # those to avoid duplicate injections.
-            if instr.lineno == lineno:
-                locs.appendleft((i, instr.name))
-                instrs.add(instr.name)
-        except AttributeError:
+        if not isinstance(instr, Instr):
             # pseudo-instruction (e.g. label)
-            pass
+            continue
+        if instr.lineno == last_lineno:
+            continue
+        last_lineno = instr.lineno
+        # Some lines might be implemented across multiple instruction
+        # offsets, and sometimes a NOP is used as a placeholder. We skip
+        # those to avoid duplicate injections.
+        if instr.lineno == lineno:
+            locs.appendleft((i, instr.name))
+            instrs.add(instr.name)
 
     if not locs:
         raise InvalidLine("Line %d does not exist or is either blank or a comment" % lineno)
@@ -119,8 +119,8 @@ def _inject_hook(code: Bytecode, hook: HookType, lineno: int, arg: Any) -> None:
         # just a placeholder.
         locs = deque((i, instr) for i, instr in locs if instr != "NOP")
 
-    for i, instr in locs:
-        if instr.startswith("END_"):
+    for i, opname in locs:
+        if opname.startswith("END_"):
             # This is the end of a block, e.g. a for loop. We have already
             # instrumented the block on entry, so we skip instrumenting the
             # end as well.
@@ -140,20 +140,26 @@ def _eject_hook(code: Bytecode, hook: HookType, line: int, arg: Any) -> None:
     """
     locs: deque[int] = deque()
     for i, instr in enumerate(code):
+        if not isinstance(instr, Instr):
+            # pseudo-instruction (e.g. label)
+            continue
         try:
             # DEV: We look at the expected opcode pattern to match the injected
             # hook and we also test for the expected opcode arguments
+            _hook_instr = code[i + _INJECT_HOOK_OPCODE_POS]
+            _arg_instr = code[i + _INJECT_ARG_OPCODE_POS]
+            _window = [code[_] for _ in range(i, i + len(_INJECT_HOOK_OPCODES))]
             if (
                 instr.lineno == line
-                and code[i + _INJECT_HOOK_OPCODE_POS].arg == hook  # bound methods don't like identity comparisons
-                and code[i + _INJECT_ARG_OPCODE_POS].arg is arg
-                and [code[_].name for _ in range(i, i + len(_INJECT_HOOK_OPCODES))] == _INJECT_HOOK_OPCODES
+                and isinstance(_hook_instr, Instr)
+                and _hook_instr.arg == hook  # bound methods don't like identity comparisons
+                and isinstance(_arg_instr, Instr)
+                and _arg_instr.arg is arg
+                and all(isinstance(_c, Instr) for _c in _window)
+                and [_c.name for _c in _window if isinstance(_c, Instr)] == _INJECT_HOOK_OPCODES
             ):
                 locs.appendleft(i)
-        except AttributeError:
-            # pseudo-instruction (e.g. label)
-            pass
-        except IndexError:
+        except (AttributeError, IndexError):
             pass
 
     if not locs:

ddtrace/internal/wrapping/__init__.py

@@ -264,7 +264,7 @@ def wrap_bytecode(wrapper: Wrapper, wrapped: FunctionType) -> bc.Bytecode:
     # Include code for handling free/cell variables, if needed
     if PY >= (3, 11):
         if code.co_cellvars:
-            instrs[0:0] = [Instr("MAKE_CELL", bc.CellVar(_), lineno=lineno) for _ in code.co_cellvars]
+            instrs[0:0] = [Instr("MAKE_CELL", bc.CellVar(_), lineno=lineno) for _ in code.co_cellvars]  # type: ignore[attr-defined]
 
         if code.co_freevars:
             instrs.insert(0, Instr("COPY_FREE_VARS", len(code.co_freevars), lineno=lineno))
@@ -316,15 +316,15 @@ def wrap(f: FunctionType, wrapper: Wrapper) -> WrappedFunction:
 
     # Copy over the code attributes
     wrapped_code.argcount = argcount
-    wrapped_code.argnames = code.co_varnames[:nargs]
+    wrapped_code.argnames = list(code.co_varnames[:nargs])
     wrapped_code.filename = code.co_filename
-    wrapped_code.freevars = code.co_freevars
-    wrapped_code.flags = flags
+    wrapped_code.freevars = list(code.co_freevars)
+    wrapped_code.flags = bc.CompilerFlags(flags)
     wrapped_code.kwonlyargcount = kwonlycount
     wrapped_code.name = code.co_name
     wrapped_code.posonlyargcount = code.co_posonlyargcount
     if PY >= (3, 11):
-        wrapped_code.cellvars = code.co_cellvars
+        wrapped_code.cellvars = list(code.co_cellvars)
 
     # Replace the function code with the trampoline bytecode
     f.__code__ = wrapped_code.to_code()

ddtrace/internal/wrapping/asyncs.py

@@ -711,7 +711,7 @@
     raise RuntimeError(msg)
 
 
-def wrap_async(instrs: list[bc.Instr], code: CodeType, lineno: int) -> None:
+def wrap_async(instrs: bc.Bytecode, 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(lineno=lineno)

ddtrace/internal/wrapping/context.py

@@ -294,7 +294,9 @@ class BaseWrappingContext(ABC):
 
     def __init__(self, f: FunctionType):
         self.__wrapped__ = f
-        self._storage: ContextVar[t.Optional[dict]] = ContextVar(f"{type(self).__name__}__storage", default=None)
+        self._storage: ContextVar[t.Optional[dict[str, t.Any]]] = ContextVar(
+            f"{type(self).__name__}__storage", default=None
+        )
 
     def __getstate__(self) -> dict[str, t.Any]:
         state = self.__dict__.copy()
@@ -315,7 +317,7 @@ def __enter__(self) -> "BaseWrappingContext":
         return self
 
     def _pop_storage(self) -> dict[str, t.Any]:
-        storage = t.cast(dict, self._storage.get())
+        storage = t.cast(dict[str, t.Any], self._storage.get())
         self._storage.set(storage.pop("__dd_wrapping_context_prev__"))
         return storage
 
@@ -332,10 +334,10 @@ def __exit__(
         self._pop_storage()
 
     def get(self, key: str) -> t.Any:
-        return t.cast(dict, self._storage.get())[key]
+        return t.cast(dict[str, t.Any], self._storage.get())[key]
 
     def set(self, key: str, value: T) -> T:
-        t.cast(dict, self._storage.get())[key] = value
+        t.cast(dict[str, t.Any], self._storage.get())[key] = value
         return value
 
     @classmethod
@@ -368,7 +370,7 @@ class WrappingContext(BaseWrappingContext):
     @property
     def __frame__(self) -> FrameType:
         try:
-            return _UniversalWrappingContext.extract(self.__wrapped__).get("__frame__")
+            return t.cast(FrameType, _UniversalWrappingContext.extract(self.__wrapped__).get("__frame__"))
         except ValueError:
             raise AttributeError("Wrapping context not entered")
 
@@ -437,7 +439,7 @@ def wrap(self) -> None:
                 super().wrap()
                 return
 
-            def trampoline(_: t.Any, args: tuple, kwargs: dict) -> t.Any:
+            def trampoline(_: t.Any, args: tuple[t.Any, ...], kwargs: dict[str, t.Any]) -> t.Any:
                 with tl:
                     f = t.cast(WrappedFunction, self.__wrapped__)
                     if is_wrapped_with(self.__wrapped__, trampoline):
@@ -526,7 +528,7 @@ def unregister(self, context: WrappingContext) -> None:
             self.unwrap()
 
     def is_registered(self, context: WrappingContext) -> bool:
-        return type(context) in self._contexts
+        return any(isinstance(c, type(context)) for c in self._contexts)
 
     def registered(self, context_type: type[WrappingContext]) -> WrappingContext:
         for context in self._contexts:
@@ -700,12 +702,8 @@ def unwrap(self) -> None:
             # Remove the head of the try block
             wc = wrapped.__dd_context_wrapped__
             for i, instr in enumerate(bc):
-                try:
-                    if instr.name == "LOAD_CONST" and instr.arg is wc:
-                        break
-                except AttributeError:
-                    # Not an instruction
-                    pass
+                if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc:
+                    break
 
             # Search for the RESUME instruction
             for i, instr in enumerate(bc, 1):
@@ -760,29 +758,19 @@ def wrap(self) -> None:
             i = 0
             while i < len(bc):
                 instr = bc[i]
-                try:
+                if isinstance(instr, bytecode.Instr):
                     if instr.name == "RETURN_VALUE":
                         return_code = CONTEXT_RETURN.bind({"context": self}, lineno=instr.lineno)
-                    else:
-                        return_code = []
-
-                    bc[i:i] = return_code
-                    i += len(return_code)
-                except AttributeError:
-                    # Not an instruction
-                    pass
+                        bc[i:i] = return_code
+                        i += len(return_code)
                 i += 1
 
             # Search for the GEN_START instruction, which needs to stay on top.
             i = 0
             if sys.version_info >= (3, 10) and (iscoroutinefunction(f) or isgeneratorfunction(f)):
                 for i, instr in enumerate(bc, 1):
-                    try:
-                        if instr.name == "GEN_START":
-                            break
-                    except AttributeError:
-                        # Not an instruction
-                        pass
+                    if isinstance(instr, bytecode.Instr) and instr.name == "GEN_START":
+                        break
 
             *bc[i:i], except_label = CONTEXT_HEAD.bind({"context": self}, lineno=code.co_firstlineno)
 
@@ -815,26 +803,18 @@ def unwrap(self) -> None:
             # Remove the head of the try block
             wc = wrapped.__dd_context_wrapped__
             for i, instr in enumerate(bc):
-                try:
-                    if instr.name == "LOAD_CONST" and instr.arg is wc:
-                        break
-                except AttributeError:
-                    # Not an instruction
-                    pass
+                if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc:
+                    break
 
             bc[i : i + len(CONTEXT_HEAD) - 1] = []
 
             # Remove all the return handlers
             i = 0
             while i < len(bc):
                 instr = bc[i]
-                try:
-                    if instr.name == "RETURN_VALUE":
-                        bc[i - len(CONTEXT_RETURN) : i] = []
-                        i -= len(CONTEXT_RETURN)
-                except AttributeError:
-                    # Not an instruction
-                    pass
+                if isinstance(instr, bytecode.Instr) and instr.name == "RETURN_VALUE":
+                    bc[i - len(CONTEXT_RETURN) : i] = []
+                    i -= len(CONTEXT_RETURN)
                 i += 1
 
             # Recreate the code object

ddtrace/internal/wrapping/generators.py

@@ -484,7 +484,7 @@
     raise RuntimeError(msg)
 
 
-def wrap_generator(instrs: list[bc.Instr], code: CodeType, lineno: int) -> None:
+def wrap_generator(instrs: bc.Bytecode, code: CodeType, lineno: int) -> None:
     if GENERATOR_HEAD_ASSEMBLY is not None:
         instrs[0:0] = GENERATOR_HEAD_ASSEMBLY.bind(lineno=lineno)
 

mypy.ini

@@ -1181,17 +1181,10 @@ disallow_untyped_calls = false
 disallow_untyped_defs = false
 disallow_incomplete_defs = false
 
-[mypy-ddtrace.internal.assembly]
-disallow_subclassing_any = false
-warn_return_any = false
-
 [mypy-ddtrace.internal.atexit]
 disallow_any_generics = false
 disallow_untyped_calls = false
 
-[mypy-ddtrace.internal.bytecode_injection.core]
-check_untyped_defs = false
-
 [mypy-ddtrace.internal.ci_visibility._api_client]
 disallow_any_generics = false
 disallow_untyped_calls = false
@@ -1944,11 +1937,6 @@ disallow_any_generics = false
 disallow_untyped_defs = false
 disallow_incomplete_defs = false
 
-[mypy-ddtrace.internal.wrapping.context]
-disallow_any_generics = false
-warn_return_any = false
-strict_equality = false
-
 [mypy-ddtrace.internal.writer.writer]
 disallow_any_generics = false
 disallow_untyped_calls = false

pyproject.toml

@@ -109,6 +109,7 @@ lint = [
     "cmake-format==0.6.13",
     "ruamel.yaml==0.18.6",
     "ast-grep-cli==0.39.4",
+    "bytecode==0.17.0",
 ]
 clean = [
     "cython",