Inject a function call into an existing function (like a breakpoint)

hook.py

from codetransformer import CodeTransformer, pattern
from codetransformer.patterns import matchany, var, option
from codetransformer.instructions import (
    CALL_FUNCTION,
    LOAD_CONST,
    NOP,
    POP_TOP,
    POP_JUMP_IF_FALSE,
)
from codetransformer.utils.instance import instance


class _Hook(CodeTransformer):
    """CodeTrasformer for adding extra function calls before a specific line.

    Parameters
    ----------
    lno : int
        The line to hook on or after.
    hook_f : callable[(mapping, mapping) -> None] or callable[() -> None]
        The callback to attach on the given line in ``f``. If
        ``with_state`` is True then this will be passed the frame locals and
        globals.
    with_state : bool
        Should the function be passed the locals and globals.
    toggle : Toggle
        The toggle object to use to enable or disable this hook.
    globals_ : dict
        The globals for the function being hooked.

    See Also
    --------
    hook
    _Strip
    """
    def __init__(self, lno, hook_f, with_state, toggle, globals_):
        super().__init__()
        self._lno = lno
        self._hook_f = hook_f
        self._with_state = with_state
        self._toggle = toggle
        self._globals = globals_
        self._hooked = False

    @pattern(...)
    def _(self, instr):
        if not self._hooked and self.code.lno_of_instr[instr] >= self._lno:
            # We are on or past the line that we want to hook and we have
            # not yet injected our call.

            # Mark that we have hooked this function.
            self._hooked = True

            # Add a marker so we can strip this out later.
            yield NOP()

            yield LOAD_CONST(self._toggle)
            # TOS  self._toggle
            yield POP_JUMP_IF_FALSE(instr)
            # empty
            # This checks the toggle to only execute if this hook is enabled.

            yield LOAD_CONST(self._hook_f)
            # TOS  = self._hook_f

            if self._with_state:
                # We are passing state to the function so we must get the
                # locals and globals.
                yield LOAD_CONST(locals)
                # TOS  = locals
                # TOS1 = self._hook_f
                yield CALL_FUNCTION(0)
                # TOS  = locals()
                # TOS1 = self._hook_f
                yield LOAD_CONST(self._globals)
                # TOS  = self._globals
                # TOS1 = locals()
                # TOS2 = self._hook_f
                yield CALL_FUNCTION(2)
                # TOS  = self._hook_f(locals(), self._globals)
            else:
                yield CALL_FUNCTION(0)
                # TOS  = self._hook_f()

            yield POP_TOP()
            # empty

        yield instr


@instance
class strip(CodeTransformer):
    """CodeTransformer that strips out instructions added with ``hook``.

    See Also
    --------
    hook
    """
    @pattern(NOP, matchany[var][option], POP_TOP)
    def _(self, *_):
        # Throw away everything between NOP and POP_TOP non-greedy.
        return iter(())


class Toggle:
    """A box around a boolean.

    Parameters
    ----------
    state : bool, optional
        The starting state of the toggle.
    """
    def __init__(self, state=True):
        self._state = state

    def set(self, state):
        self._state = state
        return state

    def enable(self):
        return self.set(True)

    def disable(self):
        return self.set(False)

    def __bool__(self):
        return self._state


def hook(f, lno, hook_f, with_state=False, inplace=False):
    """Add a hook function before a given line in a function

    Parameters
    ----------
    f : function
        The function to hook.
    lno : int
        The line to hook on or after.
    hook_f : callable[(mapping, mapping) -> None] or callable[() -> None]
        The callback to attach on the given line in ``f``. If
        ``with_state`` is True then this will be passed the frame locals and
        globals.
    with_state : bool, optional
        Should the function be passed the locals and globals.
    inplace : bool, optional
        Should this update ``f`` in place or return a new function.

    Returns
    -------
    hooked : function
        The hooked version of ``f``. If ``inplace=True``, then
        ``hooked is f``.
    toggle : Toggle
        An object that can enable or disable the hook from executing.

    See Also
    --------
    strip
    """
    toggle = Toggle()
    hooked = _Hook(lno, hook_f, with_state, toggle, f.__globals__)(f)
    if inplace:
        f.__code__ = hooked.__code__
    return hooked, toggle