Python Forum

Full Version: A modernized magic trick
You're currently viewing a stripped down version of our content. View the full version with proper formatting.
A few years ago, it occurred to me that if a Python instance obj had attributes 'spam' and 'eggs', I'd like to be able to pull these attributes with a syntax such as
spam, eggs = attrof(obj)
Unfortunately, it is not possible, because the attrof() function cannot access the attribute names written on the left hand side of the assignment statement. To extract attributes, their names must be passed as arguments of the function.

The only way to obtain such a syntax is that the attrof function accesses the position of the assignment statement where it is called in the source code or the bytecode in order to extract the names of the attributes on the left hand side.

This is implemented in the below module, which provides a decorator @use_assignment_names that allow functions to use simple names on the left of assignment statements. These functions must return exactly the same number of arguments that they found on the left-hand side.

The mechanism is based on CPython bytecode, tested with 3.10. This is a modernized version of a module that I wrote initialy for Python 2, and the bytecode has changed since then, so it won't work with old versions of Python, and it may need to be adapted to future versions of Python, so use this at your own risk.
 #!/usr/bin/env python
# module use_assignment_names
# SPDX-FileCopyrightText: 2023 Eric Ringeisen
# SPDX-License-Identifier: MIT

import functools
import opcode
import sys

__version__ = "2023.08.31"


class VarnamesError(Exception):
    pass


vars().update(
    {
        s: opcode.opmap[s]
        for s in ("CALL_FUNCTION", "UNPACK_SEQUENCE", "STORE_FAST", "STORE_NAME")
    }
)
errmsg = "simple assignment syntax 'x, y, z = ...' expected"


def _assignment_varnames(code, lasti):
    """Extract variable names from a statement of the form
    x, y, z = function(...)
    in a code objet @code where @lasti is the index of the
    CPython bytecode instruction where the function is called.

    Tested with CPython 3.10.12
    """

    delta = 2
    co = code.co_code
    i, k = lasti, co[lasti]
    if k != CALL_FUNCTION:
        raise VarnamesError(errmsg)
    i += delta
    k = co[i]
    if k == UNPACK_SEQUENCE:
        nvars = co[i + 1]
        i += delta
    else:
        nvars = int(k in (STORE_FAST, STORE_NAME))
    for _ in range(nvars):
        k, oparg = co[i], co[i + 1]
        if k == STORE_FAST:
            yield code.co_varnames[oparg]
        elif k == STORE_NAME:
            yield code.co_names[oparg]
        else:
            raise VarnamesError(errmsg)
        i += delta


def use_assignment_names(func):
    """use_assignment_names(function) -> decorated function

    This decorator allows a function to extract variable names
    from the line of code where it is called, and create
    values for these variables which depend on their names.

    Argument:
        function : a function having with a variadic argument *names
                and returning a sequence of the same length.
                It may have positional arguments before *names and
                keyword arguments after *names.

    WARNING: This code relies on the structure of the bytecode produced
        by CPython's interpreter. It is tested with Python 3.10. It is
        known to fail for old versions of Python.

    Usage:

        The decorated functions can be used in simple assignment expressions:

            @use_assignment_names
            def func(*names):
                ...

            a, b, c = func()

        Only one-word variable names are allowed on the left-hand side.

    example:

        >>> from importlib import import_module
        >>>
        >>> @use_assignment_names
        ... def mod_pull(modname, *names):
        ...     mod = import_module(modname)
        ...     for name in names:
        ...         yield getattr(os, name)
        ...
        >>> close, closerange = mod_pull('os')
        >>> close, closerange
        <built-in function close> <built-in function closerange>
    """

    @functools.wraps(func)
    def wrapper(*args, **kwd):
        f = sys._getframe(1)
        try:
            code, lasti = f.f_code, f.f_lasti
        finally:
            del f
        more = list(_assignment_varnames(code, lasti))
        seq = func(*args, *more, **kwd)
        return next(iter(seq), None) if len(more) <= 1 else seq

    return wrapper


@use_assignment_names
def attrof(obj, *names):
    """Pull an object's attributes in an assignment statement

    Usage:

        spam, eggs = attrof(some_object)
    """
    return (getattr(obj, name) for name in names)


if __name__ == "__main__":
    from importlib import import_module

    @use_assignment_names
    def mod_pull(modname, *names):
        mod = import_module(modname)
        for name in names:
            yield getattr(mod, name)

    close, closerange = mod_pull("os")
    print(close, closerange)