Стиль удава: компиляция на лету

Это первая статья серии посвященной Кунг-фу Python. При неумелом использовании описываемых приемов можно запросто заехать себе пяткой в лоб, но настоящему программисту и это должно прийтись по душе, т. к. в конечном счете и такой опыт засчитывается в плюс.

Я решил начать с чего попроще — runtime компиляции кода. Мы не будем писать JIT-компилятор в обычном понимании, а реализуем небольшую оптимизацию для вполне практичной библиотеки.

Контекст задачи

Для целей этой статьи будем считать, что нам позарез понадобилось снабжать наши функции проверкой типов аргументов. На самом деле, по ходу дела будет видно, что нетрудно расширить библиотеку для более полезных задач как приведение аргументов к типу или любые другие произвольные преобразования. Но сначала просто проверка типов.

@accepts(float, float)
def float_div(a, b):
    return a / b

Вот так мы будем её использовать. Это иллюстрация, сейчас нам главное саму библиотеку написать.

Какой эффект мы ожидаем от такого декоратора? Давайте сразу на коленке добавим документацию и тесты, сразу и прояснится:

@accepts(float, float)
def float_div(a, b):
    """
    A function that divides floats. 
    Any other types of arguments will throw an exception

        >>> float_div(1.0, 2.0)
        0.5
        >>> float_div(1, 1.0)
        Traceback (most recent call last):
        ...
        TypeError: Expected <type 'float'> for argument #1 got <type 'int'>
    """
    return a / b


if __name__ == '__main__':
    import doctest
    doctest.run_docstring_examples(float_div, globals())

Теперь должно стать очевидно, как реализовывать декоратор:

class accepts_wrap(object):
    def __init__(self, func, argtypes):
        self.func = func
        self.__doc__ = func.__doc__
        self.argtypes = argtypes

    def __call__(self, *args):
        for i, (arg, cls) in enumerate(zip(args, self.argtypes)):
            if not isinstance(arg, cls):
                raise TypeError("Expected %r for argument #%d got %r" % (cls, i+1, type(arg)))
        return self.func(*args)


def accepts(*argtypes):
    return lambda func: accepts_wrap(func, argtypes)

Постановка задачи

Хоть ничего страшного в дополнительном коде нет, такие декораторы могут быть применены к весьма часто вызываемым функциям, и потому очень хотелось бы минимизировать возможные потери скорости. Это может быть оправдано, если мы действительно заметили замедление или если мы собираемся положить нашу библиотеку в основу многих разработок и суммарная экономия вычислительной мощности в итоге оправдает затраты времени на оптимизацию.

Также, будучи честным с самим собой, стоит признать, что писать такие нетривиальные вещи довольно интересно и потому, при наличии времени, можно реализовать что-нибудь эдакое в качестве практики, которая вполне может сильно пригодиться позже, когда оптимизация действительно понадобится и времени на усвоение приемов и ловушек такого подхода будет маловато. По сути, при должных навыках эти наши эксперименты можно переложить на генерацию машинного кода, что может очень пригодиться в определенных областях.

Так вот, мы попробуем минимизировать количество дополнительных команд исполняемых при каждом вызове обернутой функции путем генерации оптимального байт-кода для каждого экземпляра обертки.

Модель исполнения

Поскольку в данном случае нас волнует только CPython и его байт-код, то кратко рассмотрим как он работает. При первой загрузке модуля его исходники разбираются в AST (Abstract Syntax Tree) по которому генерируется байт-код (который также кэшируется в.pyc файле). Байт-код python предназначен для исполнения несложной умеренно сложной стековой виртуальной машиной. Реализация внутреннего цикла CPython находится в файле ceval.c и ознакомиться с ним хотя бы поверхностно следует всякому интересующемуся тем, как Python работает внутри. Если вы собираетесь писать компиляцию на лету как в этой статье, то такое знакомство, хотя и не обязательно, но придется очень кстати.

Инструментарий

Python предоставляет доступ к разным этапам процесса загрузки и генерации байт-кода начиная с разбития на лексемы и построения синтаксических деревьев и заканчивая созданием из этих структур code-objects и собственно функций, классов и модулей, которые можно использовать непосредственно. Есть также возможность загружать и сохранять байткод самостоятельно, т. к. формат.pyc-файлов это знакомый всем marshal. Еще очень полезен в разработке декомпилятор байт-кода, позволяющий посмотреть какие коды получаются из существующих исходников.

В зависимости от стоящих задач можно подойти к runtime-компиляции по-разному. Ведь если для AST / compileast можно ожидать поддержки в разных реализациях интерпретатора, то для собственно байт-кода этого сказать нельзя. Так что, возможно, для многих лучшим выбором окажется работа на другом уровне, однако в этой статье мы будем генерировать байт-код непосредственно. Для этого мы воспользуемся библиотекой BytecodeAssembler.

Порождающий код, написанный с её использованием очень похож на ассемблер, но нужно понимать, что на самом деле это просто построение последовательности операций и четко различать какие операции выполняются, а какие добавляются в порожденный код. Это несложно, но требует некоторой адаптации. Например, вот как выглядит дизассемблированный код простенькой функции:

>>> def foo(x, y):
...     return x + y
...
>>> from dis import dis
>>> dis(foo)
  2           0 LOAD_FAST                0 (x)
              3 LOAD_FAST                1 (y)
              6 BINARY_ADD
              7 RETURN_VALUE

А вот как можно породить такой же такой код:

from peak.util.assembler import *
c = Code()
c.LOAD_FAST('x')
c.LOAD_FAST('y')
c.BINARY_ADD()
c.RETURN_VALUE()
c.co_argcount = 2
c.co_varnames = ['x', 'y']
foo_code = c.code()

Новая реализация

Дизассемблировав внутренний цикл нашей первой реализации accepts_wrap, получим кальку, по которой можно записать новый вариант:

from peak.util.assembler import *
import new
import inspect

class accepts_wrap(object):
    def __init__(self, func, argtypes):
        self.func = func
        self.__doc__ = func.__doc__
        self.argtypes = argtypes
        self.compile()

    def compile(self):
        argnames = inspect.getargspec(self.func)[0]
        c = Code.from_function(self.func)
        c.LOAD_CONST(self.func)
        for i, cls in enumerate(self.argtypes):
            c.LOAD_FAST(argnames[i])
            c.LOAD_CONST(isinstance) # if not isinstance(arg, cls):
            c.LOAD_FAST(argnames[i])
            c.LOAD_CONST(cls)
            c.CALL_FUNCTION(2)

            skip = c.JUMP_IF_TRUE() 
            c.LOAD_CONST(TypeError) # raise TypeError
            c.LOAD_CONST("Expected %r for argument #%d got %%r" % (cls, i+1))
            c.LOAD_CONST(type)
            c.LOAD_FAST(argnames[i])
            c.CALL_FUNCTION(1)
            c.BINARY_MODULO()
            c.CALL_FUNCTION(1)
            c.RAISE_VARARGS(1)

            skip()
            c.POP_TOP()
        c.CALL_FUNCTION(len(argnames))
        c.RETURN_VALUE()
        self.wrapped_func = new.function(c.code(), {})

    def __call__(self, *args):
        return self.wrapped_func(*args)

Не будем вдаваться в подробности лишь отметив, что многие операции были заменены на более дешевые. Например, isinstance загружается через LOAD_CONST, а не LOAD_GLOBAL как в первом варианте. Также исчезла итерация по списку ожидаемых классов — при генерации байт-кода цикл разворачивается.

Для того чтобы понять, как всё работает, нужно иметь четкую ментальную модель стековой машины, потому что иначе загвоздки будут возникать на каждом шагу. Например, почему мы в первую очередь делаем c. LOAD_CONST (self.func)?

После вызова isinstance на вершине стека оказывается булевое значение, по которому мы проверяем валидность аргумента. В байт-коде ветвление, естественно, реализуется через JUMP. В BytecodeAssembler для прыжков на еще не объявленные лейблы есть удобное соглашение (см. в коде: skip =... создает отложенный лейбл, а skip () привязывает его к месту в коде).

Я думаю, вы заметили, что часть с созданием экземпляра TypeError оказалась громоздкой, она повторяет достаточно прямолинейный код, не требующий к тому же никакой оптимизации, нельзя ли записать это короче? Можно. При помощи конструкций в какой-то мере напоминающих AST, вот чем можно заменить этот код:

c.LOAD_CONST(TypeError)
c.LOAD_CONST("Expected %r for argument #%d got %%r" % (cls, i+1))
c(Call(Const(type), [Local(argnames[i])]))
c.BINARY_MODULO()
c.CALL_FUNCTION(1)
c.RAISE_VARARGS(1)

Еще одна итерация и:

error_msg = Suite([ Const("Expected %r for argument #%d got %%r" % (cls, i+1)),
                    Call(Const(type), [Local(argnames[i])]),
                    Code.BINARY_MODULO])
c.LOAD_CONST(TypeError)
c(error_msg)
c.CALL_FUNCTION(1)
c.RAISE_VARARGS(1)

И еще немного:

error_msg = Suite([ Const("Expected %r for argument #%d got %%r" % (cls, i+1)),
                    Call(Const(type), [Local(argnames[i])]),
                    Code.BINARY_MODULO])
c(Call(Const(TypeError), [error_msg]))
c.RAISE_VARARGS(1)

Проделав аналогичные итерации с остальным циклом, получим следующее:

        for i, cls in enumerate(self.argtypes):
            arg = Local(argnames[i])
            c(arg)
            c(Call(Const(isinstance), [arg, Const(cls)]))
            skip = c.JUMP_IF_TRUE()
            error_msg = Suite([ Const("Expected %r for argument #%d got %%r" % (cls, i+1)),
                                Call(Const(type), [arg]),
                                Code.BINARY_MODULO])
            c(Call(Const(TypeError), [error_msg]))
            c.RAISE_VARARGS(1)
            skip()
            c.POP_TOP()

Расширяемость

Далее, чтобы добиться расширяемости, реализуем тело цикла как новую конструкцию.

Для реализации process_arg есть вспомогательные средства, с ними он будет выглядеть следующим образом:

@nodetype()
def process_arg(i, argname, spec, code=None):
    if code is None:
        return i, argname, spec
    else:
        arg = Local(argname)
        code(Call(Const(isinstance), [arg, Const(spec)]))
        skip = code.JUMP_IF_TRUE()
        error_msg = Suite([ Const("Expected %r for argument #%d got %%r" % (spec, i+1)),
                            Call(Const(type), [arg]),
                            Code.BINARY_MODULO])
        code(Call(Const(TypeError), [error_msg]))
        code.RAISE_VARARGS(1)
        skip()
        code.POP_TOP()
        code(arg)

Это в свою очередь позволяет нам сократить реализацию compile:

    def compile(self):
        argnames = inspect.getargspec(self.func)[0]
        c = Code.from_function(self.func)
        code_args = [process_arg(i, argnames[i], cls) for i, cls in enumerate(self.argtypes)]
        c.return_(Call(Const(self.func), code_args))
        self.wrapped_func = new.function(c.code(), {})

На самом деле можно сократить весь класс accepts_wrap, вот что от него осталось:

def accepts(*argtypes):
    def decorate(func):
        argnames = inspect.getargspec(func)[0]
        code_args = [process_arg(i, argnames[i], cls) for i, cls in enumerate(argtypes)]
        c = Code.from_function(func)
        c.return_(Call(Const(func), code_args))
        wrapped = new.function(c.code(), {})
        return rewrap(func, wrapped) # preserve name / __doc__
    return decorate

Вот весь наш код на данный момент, совсем немного. Благодаря тому, что обработка каждого аргумента вынесена в process_arg, несложно переопределять её для особых случаев. Например, на данный момент мы предполагаем что аргументами декоратора будут классы, можно же предположить, что ими могут также быть некие другие объекты, описывающие требования к аргументу. При условии, что мы умеем генерировать для этих проверок оптимальный код (эту задачу можно делегировать самим декларациям) мы получим свободно расширяемый предобработчкик аргументов. Естественно, в реальном проекте начинать надо с более традиционной реализации, но данная статья, я надеюсь, в какой-то мере демонстрирует каким образом возможно усиливать некоторые библиотеки частичной компиляцией на лету.

Само собой без замеров относительной скорости оптимизировать реализацию не рекомендуется. Также следует помнить, что генерация кода доступна не во всех окружениях и для максимальной совместимости нужно также поддерживать и тестировать обычную реализацию той же функциональности.

Интересное про байт-код и интерпретатор CPython

Для того чтобы комфортно писать генерацию байт-кода нужно ознакомиться с внутренностями интерпретатора, но поскольку из читателей наверняка этим займутся лишь немногие, перескажу несколько интересных фактов.

• Для получения атрибута нет специальной команды, компилятор генерирует последовательность команд LOAD_CONST (getattr), LOAD_XXX (obj), LOAD_CONST (attrname), CALL_FUNCTION (2)
• Локальные переменные во много раз быстрее глобальных. Поиск глобальной переменной это доступ к словарю, а для локальной переменной это доступ к ячейке стека по заранее известному индексу.
• Переменные общие для замыканий или объявленные уровнем выше реализованы через специальные структуры данных, для доступа к ним существуют отдельные команды.
• На уровне интерпретатора реализовано предсказание последовательностей команд для ускорения наиболее часто встречающихся цепочек.

Если есть какие-то вопросы, связанные с этой темой, спрашивайте в комментариях, попробую ответить.

Підписуйтеся на Telegram-канал «DOU #tech», щоб не пропустити нові технічні статті.