_trace_func

COMPARE_OPS: Dict[str, Callable] = {'<': operator.lt, '<=': operator.le, '==': operator.eq, '!=': operator.neg, '>': operator.gt, '>=': operator.ge, 'in': operator.contains, 'is': operator.is_, 'is not': operator.is_not, 'not in': composer(operator.not_, operator.contains)} module-attribute

Defer supporting imports until after imports are turned into call nodes Usually we don't need to worry about it, at least in the context of slicing, but consider the following example:

import x
if ...:
  import x.y
a = x.y.z + 1

lineapy.save(a, "weird value") actually would have a dependency on the blackbox, as discussed in https://github.com/LineaLabs/lineapy/blob/main/docs/source/rfcs/0001-imports.md. Can read more: https://docs.python.org/3/reference/import.html#submodules

Here are the ones we don't support:

• IMPORT_STARIMPORT_STAR
• IMPORT_NAME
• IMPORT_FROM

all_code_objects(code)

Return all code objects from a source code object. This will include those used within it, such as nested functions.

Source code in lineapy/system_tracing/_trace_func.py

def all_code_objects(code: CodeType) -> Iterable[CodeType]:
    """
    Return all code objects from a source code object. This will include those used within it, such as nested functions.
    """
    yield code
    for const in code.co_consts:
        if isinstance(const, CodeType):
            yield from all_code_objects(const)

resolve_bytecode_execution(name, value, arg, stack, offset, extended_arg_count)

Returns a function call corresponding to the bytecode executing on the current stack.

Source code in lineapy/system_tracing/_trace_func.py

def resolve_bytecode_execution(
    name: str,
    value: Any,
    arg: Optional[int],
    stack: OpStack,
    offset: int,
    extended_arg_count: int,
) -> Union[Iterable[FunctionCall], ReturnValueCallback, FunctionCall, None]:
    """
    Returns a function call corresponding to the bytecode executing on the current stack.
    """
    if name in NOT_FUNCTION_CALLS:
        return None
    if name in UNARY_OPERATORS:
        # Unary operations take the top of the stack, apply the operation, and push the
        # result back on the stack.
        args = [stack[-1]]
        return lambda post_stack, _: FunctionCall(
            # updating mypy to 0.942 causes issues here. mypy 0.931 works fine.
            # something is up with the iter because the same thing works for binary ops below
            UNARY_OPERATORS[name],  # type: ignore
            args,
            {},
            post_stack[-1],
        )
    if name in BINARY_OPERATIONS:
        # Binary operations remove the top of the stack (TOS) and the second top-most
        # stack item (TOS1) from the stack.  They perform the operation, and put the
        # result back on the stack.
        args = [stack[-2], stack[-1]]
        return lambda post_stack, _: FunctionCall(
            BINARY_OPERATIONS[name], args, {}, post_stack[-1]
        )
    if name == "FOR_ITER":
        """
        From the original Python source code documentation:
        # TOS is an `iterator`.  Call its `__next__` method.  If
        # this yields a new value, push it on the stack (leaving the iterator below
        # it).  If the iterator indicates it is exhausted, TOS is popped, and the byte
        # code counter is incremented by *delta*.

        We need to handle the case when the code body is large and needs
        `EXTENDED_ARG`, this breaks the assumption that the offset will increase
        by 2, so we have the count of extended_arg passed in as traced by the caller
        and count with that instead.

        ```
        >>> large_for_loop_code = "for _ in x:\n  i = 1\n" + "  j = i\n" * 100
        >>> dis.dis(large_for_loop_code)
        1           0 LOAD_NAME                0 (x)
                    2 GET_ITER
                    4 EXTENDED_ARG             1
                    6 FOR_ITER               408 (to 416)
                    8 STORE_NAME               1 (_)

        2          10 LOAD_CONST               0 (1)
        ```

        So we translated to:
        If the current instruction is the next one (i.e. the offset has increased by 2), then we didn't jump,
        meaning the iterator was not exhausted. Otherwise, we did jump, and it was, so don't add a function call for this.

        Note that if we start handling exception, we should edit how we are
        handling the loops, since Python uses the try/catch mechanism to catch
        loop end.
        """
        args = [stack[-1]]

        return (
            lambda post_stack, post_offset: FunctionCall(
                next, args, {}, post_stack[-1]
            )
            if post_offset == offset + 2 * extended_arg_count
            else None
        )
    if name == "STORE_SUBSCR":
        # Implements ``TOS1[TOS] = TOS2``.
        return FunctionCall(
            operator.setitem, [stack[-2], stack[-1], stack[-3]]
        )
    if name == "DELETE_SUBSCR":
        # Implements ``del TOS1[TOS]``.
        return FunctionCall(operator.delitem, [stack[-2], stack[-1]])
    if name == "SET_ADD":
        # Calls ``set.add(TOS1[-i], TOS)``.  Used to implement set comprehensions.
        return FunctionCall(getattr(stack[-value - 1], "add"), [stack[-1]])
    if name == "LIST_APPEND":
        # Calls `list.append(TOS1[-i], TOS)`.  Used to implement list comprehensions.
        return FunctionCall(getattr(stack[-value - 1], "append"), [stack[-1]])

    if name == "MAP_ADD":
        #  Calls `dict.__setitem__(TOS1[-i], TOS1, TOS)`.  Used to implement dict comprehensions.
        dict_ = stack[-value - 2]
        # Version 3.8:  Map value is TOS and map key is TOS1. Before, those were reversed.
        if version_info >= (3, 8):
            key = stack[-2]
            value = stack[-1]
        else:
            key = stack[-1]
            value = stack[-2]
        return FunctionCall(operator.setitem, [dict_, key, value])
    if name == "WITH_EXCEPT_START":
        # Calls the function in position 7 on the stack with the top three
        # items on the stack as arguments.
        # Used to implement the call ``context_manager.__exit__(*exc_info())`` when an exception
        # has occurred in a :keyword:`with` statement.
        fn = stack[-7]
        args = [stack[-1], stack[-2], stack[-3]]
        return lambda post_stack, _: FunctionCall(fn, args, res=post_stack[-1])

    if name == "SETUP_WITH":
        # This opcode performs several operations before a with block starts.  First,
        # it loads `__exit__` from the context manager and pushes it onto
        # the stack for later use by `WITH_EXCEPT_START`.  Then,
        # `__enter__` is called, and a finally block pointing to *delta*
        # is pushed.  Finally, the result of calling the ``__enter__()`` method is pushed onto
        # the stack.  The next opcode will either ignore it (`POP_TOP`), or
        # store it in (a) variable(s) (`STORE_FAST`, `STORE_NAME`, or
        # `UNPACK_SEQUENCE`).
        x = stack[-1]
        # The __enter__ bound method is never saved to the stack, so we recompute it to save in the function call
        return lambda post_stack, _: FunctionCall(
            getattr(x, "__enter__"), [], res=post_stack[-1]
        )

    if name == "WITH_CLEANUP_START":
        # At the top of the stack are 1-6 values indicating
        # how/why we entered the finally clause:
        # - TOP = None
        # - (TOP, SECOND) = (WHY_{RETURN,CONTINUE}), retval
        # - TOP = WHY_*; no retval below it
        # - (TOP, SECOND, THIRD) = exc_info()
        #     (FOURTH, FIFTH, SIXTH) = previous exception for EXCEPT_HANDLER
        # Below them is EXIT, the context.__exit__ bound method.
        # In the last case, we must call
        #     EXIT(TOP, SECOND, THIRD)
        # otherwise we must call
        #     EXIT(None, None, None)

        # In the first three cases, we remove EXIT from the
        # stack, leaving the rest in the same order.  In the
        # fourth case, we shift the bottom 3 values of the
        # stack down, and replace the empty spot with NULL.

        # In addition, if the stack represents an exception,
        # *and* the function call returns a 'true' value, we
        # push WHY_SILENCED onto the stack.  END_FINALLY will
        # then not re-raise the exception.  (But non-local
        # gotos should still be resumed.)
        args = [None, None, None]
        try:
            exc = stack[-1]
        except ValueError:
            exc = None
            fn = stack[-2]
        if exc is None:
            fn = stack[-2]
        elif isinstance(exc, int):
            # WHY_RETURN and WHY_CONTINUE
            if exc in {0x0008, 0x0020}:
                fn = stack[-3]
            else:
                fn = stack[-2]
        else:
            args = [exc, stack[-2], stack[-3]]
            fn = stack[-7]
        return lambda post_stack, _: FunctionCall(fn, args, res=post_stack[-1])
    if name == "UNPACK_SEQUENCE":
        # Unpacks TOS into *count* individual values, which are put onto the stack
        # right-to-left.
        from lineapy.utils.lineabuiltins import l_unpack_sequence

        seq = stack[-1]

        def callback(post_stack: OpStack, _):
            # Replicate the behavior of using our internal functions for unpacking
            unpacked = [post_stack[-i - 1] for i in range(value)]
            yield FunctionCall(l_unpack_sequence, [seq, value], res=unpacked)
            for i, v in enumerate(unpacked):
                yield FunctionCall(operator.getitem, [unpacked, i], res=v)

        return callback
    if name == "UNPACK_EX":
        # Implements assignment with a starred target: Unpacks an iterable in TOS into
        # individual values, where the total number of values can be smaller than the
        # number of items in the iterable: one of the new values will be a list of all
        # leftover items.
        #
        # The low byte of *counts* is the number of values before the list value, the
        # high byte of *counts* the number of values after it.  The resulting values
        # are put onto the stack right-to-left.

        from lineapy.utils.lineabuiltins import l_unpack_ex

        seq = stack[-1]

        count_left = cast(int, arg) % 256
        count_right = cast(int, arg) >> 8

        def callback(post_stack: OpStack, _):
            unpacked = [
                post_stack[-i - 1] for i in range(count_left + count_right + 1)
            ]
            yield FunctionCall(
                l_unpack_ex, [seq, count_left, count_right], res=unpacked
            )
            for i, v in enumerate(unpacked):
                yield FunctionCall(operator.getitem, [unpacked, i], res=v)

        return callback

    if name == "STORE_ATTR":
        # Implements ``TOS.name = TOS1``, where *namei* is the index of name in
        # `co_names`.
        return FunctionCall(setattr, [stack[-1], value, stack[-2]])
    if name == "DELETE_ATTR":
        # Implements ``del TOS.name``, using *namei* as index into :attr:`co_names`.
        return FunctionCall(delattr, [stack[-1], value])

    if name in {"BUILD_TUPLE", "BUILD_LIST", "BUILD_SET"}:
        # Creates a tuple consuming *count* items from the stack, and pushes the
        # resulting tuple onto the stack.
        from lineapy.utils.lineabuiltins import l_list, l_set, l_tuple

        INSTRUCTION_TO_FN = {
            "BUILD_TUPLE": l_tuple,
            "BUILD_LIST": l_list,
            "BUILD_SET": l_set,
        }
        fn = INSTRUCTION_TO_FN[name]

        args = [stack[-i - 1] for i in range(value)]
        args.reverse()
        return lambda post_stack, stack_offset: FunctionCall(
            fn, args, res=post_stack[-1]
        )

    if name == "LIST_TO_TUPLE":
        args = [stack[-1]]
        return lambda post_stack, stack_offset: FunctionCall(
            tuple, args, res=post_stack[-1]
        )
    # Python < 3.9
    if name == "BUILD_TUPLE_UNPACK":
        from lineapy.utils.lineabuiltins import l_list

        # Compiles down to one empty list creation, followed by a number of extends, and then convert to tuple
        args = [stack[-i - 1] for i in range(value)]
        args.reverse()

        def callback(post_stack: OpStack, _):
            intermediate_list = list(post_stack[-1])
            extend = intermediate_list.extend
            return [
                FunctionCall(l_list, res=intermediate_list),
                *(FunctionCall(extend, [a]) for a in args),
                FunctionCall(tuple, [intermediate_list], res=post_stack[-1]),
            ]

        return callback

    if name == "BUILD_LIST_UNPACK":

        # Compiles down to one empty list creation, followed by a number of extends
        args = [stack[-i - 1] for i in range(value)]
        args.reverse()
        return lambda post_stack, stack_offset: [
            FunctionCall(list, res=post_stack[-1])
        ] + [
            FunctionCall(getattr(post_stack[-1], "extend"), [a]) for a in args
        ]
    if name == "BUILD_TUPLE_UNPACK":

        # Compiles down to one empty list creation, followed by a number of extends, and then convert to tuple
        args = [stack[-i - 1] for i in range(value)]
        args.reverse()

        def callback(post_stack: OpStack, _):
            intermediate_list = list(post_stack[-1])
            extend = intermediate_list.extend
            return [
                FunctionCall(l_list, res=intermediate_list),
                *(FunctionCall(extend, [a]) for a in args),
                FunctionCall(tuple, [intermediate_list], res=post_stack[-1]),
            ]

        return callback
    if name == "BUILD_SET_UNPACK":
        from lineapy.utils.lineabuiltins import l_set

        # Compiles down to one empty set creation, followed by a number of update
        args = [stack[-i - 1] for i in range(value)]
        args.reverse()
        return lambda post_stack, stack_offset: [
            FunctionCall(l_set, res=post_stack[-1])
        ] + [
            FunctionCall(getattr(post_stack[-1], "update"), [a]) for a in args
        ]
    if name == "BUILD_MAP_UNPACK":
        from lineapy.utils.lineabuiltins import l_dict

        # Compiles down to one empty dic creation, followed by a number of extends
        args = [stack[-i - 1] for i in range(value)]
        args.reverse()
        return lambda post_stack, stack_offset: [
            FunctionCall(l_dict, res=post_stack[-1])
        ] + [
            FunctionCall(getattr(post_stack[-1], "update"), [a]) for a in args
        ]
    if name == "BUILD_MAP":
        # Pushes a new dictionary object onto the stack.  Pops ``2 * count`` items
        # so that the dictionary holds *count* entries:
        # ``{..., TOS3: TOS2, TOS1: TOS}``.
        from lineapy.utils.lineabuiltins import l_dict, l_tuple  # noqa: F811

        args = [(stack[-i * 2 - 2], stack[-i * 2 - 1]) for i in range(value)]
        args.reverse()

        return lambda post_stack, stack_offset: [
            FunctionCall(l_tuple, [k, v], res=(k, v)) for k, v in args
        ] + [FunctionCall(l_dict, args, res=post_stack[-1])]

    if name == "LIST_EXTEND":
        return FunctionCall(getattr(stack[-value - 1], "extend"), [stack[-1]])

    if name == "SET_UPDATE":
        return FunctionCall(getattr(stack[-value - 1], "update"), [stack[-1]])

    if name == "DICT_UPDATE" or name == "DICT_MERGE":
        return FunctionCall(getattr(stack[-value - 1], "update"), [stack[-1]])

    if name == "LOAD_ATTR":
        o = stack[-1]
        return lambda post_stack, _: FunctionCall(
            getattr, [o, value], res=post_stack[-1]
        )

    if name == "COMPARE_OP":
        args = [stack[-2], stack[-1]]
        if value == "in":
            args.reverse()
        return lambda post_stack, _: FunctionCall(
            COMPARE_OPS[value], args, res=post_stack[-1]
        )

    if name == "IS_OP":
        args = [stack[-2], stack[-1]]
        return lambda post_stack, _: FunctionCall(
            operator.is_not if value == 1 else operator.is_,
            args,
            res=post_stack[-1],
        )
    if name == "CONTAINS_OP":
        args = [stack[-1], stack[-2]]
        return lambda post_stack, _: FunctionCall(
            operator.contains, args, res=post_stack[-1]
        )
    if name == "CALL_FUNCTION":
        # Calls a callable object with positional arguments.
        # *argc* indicates the number of positional arguments.
        # The top of the stack contains positional arguments, with the right-most
        # argument on top.  Below the arguments is a callable object to call.
        # `CALL_FUNCTION` pops all arguments and the callable object off the stack,
        # calls the callable object with those arguments, and pushes the return value
        # returned by the callable object.
        args = list(reversed([stack[-i - 1] for i in range(value)]))
        fn = stack[-value - 1]
        return lambda post_stack, _: FunctionCall(fn, args, res=post_stack[-1])

    if name == "CALL_FUNCTION_KW":
        # Calls a callable object with positional (if any) and keyword arguments.
        # *argc* indicates the total number of positional and keyword arguments.
        # The top element on the stack contains a tuple with the names of the
        # keyword arguments, which must be strings.
        # Below that are the values for the keyword arguments,
        # in the order corresponding to the tuple.
        # Below that are positional arguments, with the right-most parameter on
        # top.  Below the arguments is a callable object to call.
        # ``CALL_FUNCTION_KW`` pops all arguments and the callable object off the stack,
        # calls the callable object with those arguments, and pushes the return value
        # returned by the callable object.
        kwarg_names: Tuple[str, ...] = stack[-1]
        n_kwargs = len(kwarg_names)
        kwargs = {
            k: stack[-i - 2] for i, k in enumerate(reversed(kwarg_names))
        }
        args = [stack[-i - 2] for i in reversed(range(n_kwargs, value))]
        fn = stack[-value - 2]
        return lambda post_stack, _: FunctionCall(
            fn, args, kwargs, post_stack[-1]
        )

    if name == "CALL_FUNCTION_EX":
        # the only case that we cannot handle is an generator that exhausts
        #   then `raise NotImplementedError()`
        # The way we figure out if something is an iterator without accidentally calling .next, is to check whether it's a Sequence, since a generator doesn't have __getitem__ (the streaming/lazy semantic) https://docs.python.org/3/library/collections.abc.html#collections-abstract-base-classes
        if cast(int, arg) & 0x01:
            # then it's kwargs
            kwargs = stack[-1]
            args = stack[-2]
            fn = stack[-3]
        else:
            # then it's positional
            args = stack[-1]
            fn = stack[-2]
        # check if the function is a generator
        if not isinstance(args, Sequence):
            raise NotImplementedError()
        return lambda post_stack, _: FunctionCall(
            fn, list(args), kwargs, post_stack[-1]
        )

    if name == "CALL_METHOD":
        # Calls a method.  *argc* is the number of positional arguments.
        # Keyword arguments are not supported.  This opcode is designed to be used
        # with `LOAD_METHOD`.  Positional arguments are on top of the stack.
        # Below them, the two items described in `LOAD_METHOD` are on the
        # stack (either ``self`` and an unbound method object or ``NULL`` and an
        # arbitrary callable). All of them are popped and the return value is pushed.

        args = list(reversed([stack[-i - 1] for i in range(value)]))
        self_ = stack[-value - 1]
        try:
            method = stack[-value - 2]
        # This is raised when it is NULL
        except ValueError:
            fn = self_
        else:
            fn = getattr(self_, method.__name__)
        return lambda post_stack, _: FunctionCall(fn, args, res=post_stack[-1])

    if name == "BUILD_SLICE":
        # Pushes a slice object on the stack.  *argc* must be 2 or 3.  If it is 2,
        # ``slice(TOS1, TOS)`` is pushed; if it is 3, ``slice(TOS2, TOS1, TOS)`` is
        # pushed. See the :func:`slice` built-in function for more information.
        if value == 2:
            args = [stack[-2], stack[-1]]
        elif value == 3:
            args = [stack[-3], stack[-2], stack[-1]]
        else:
            raise NotImplementedError()

        return lambda post_stack, _: FunctionCall(
            slice, args, res=post_stack[-1]
        )
    if name == "FORMAT_VALUE":
        # Used for implementing formatted literal strings (f-strings).  Pops
        # an optional *fmt_spec* from the stack, then a required *value*.
        # *flags* is interpreted as follows:
        #
        # * ``(flags & 0x03) == 0x00``: *value* is formatted as-is.
        # * ``(flags & 0x03) == 0x01``: call :func:`str` on *value* before
        #     formatting it.
        # * ``(flags & 0x03) == 0x02``: call :func:`repr` on *value* before
        #     formatting it.
        # * ``(flags & 0x03) == 0x03``: call :func:`ascii` on *value* before
        #     formatting it.
        # * ``(flags & 0x04) == 0x04``: pop *fmt_spec* from the stack and use
        #     it, else use an empty *fmt_spec*.
        #
        # Formatting is performed using :c:func:`PyObject_Format`.  The
        # result is pushed on the stack.
        have_format_spec = (cast(int, arg) & 0x04) == 0x04

        if have_format_spec:
            format_spec = stack[-1]
            str_ = stack[-2]
        else:
            format_spec = None
            str_ = stack[-1]

        which_conversion = cast(int, arg) & 0x03
        if which_conversion == 0x00:
            transformed_str = str_
            transform_calls = []
        elif which_conversion == 0x01:
            transformed_str = str(str_)
            transform_calls = [FunctionCall(str, [str_], res=transformed_str)]
        elif which_conversion == 0x02:
            transformed_str = repr(str_)
            transform_calls = [FunctionCall(repr, [str_], res=transformed_str)]
        elif which_conversion == 0x03:
            transformed_str = ascii(str_)
            transform_calls = [
                FunctionCall(ascii, [str_], res=transformed_str)
            ]
        else:
            raise NotImplementedError()
        return lambda post_stack, _: transform_calls + (
            [
                FunctionCall(
                    format,
                    [transformed_str, format_spec],
                    res=post_stack[-1],
                )
            ]
        )
    raise NotImplementedError()

---

Was this helpful?

Help us improve docs with your feedback!

Send