Блог - Linux, программирование, Я!

PythonСамая короткая запись асинхронных вызовов в tornado v2, или патчим Python AST

Меня очень заинтересовала статья Самая короткая запись асинхронных вызовов в tornado или патчим байткод в декораторе, не столько с практической точки зрения, сколько с точки зрения реализации.
Всё-таки модификация байткода в рантайме это слишком опасная и ненадежная операция. И уж наверняка не поддерживаемая альтернативными интерпретаторами Python.

Попробуем исправить этот недостаток способом, который для этого куда больше предназначен и который применяется для схожих целей во многих других языках (я точно встречал в Lisp и Erlang). Этот способ - модификация Абстрактного синтаксического дерева (AST) программы.

Для начала - что такое AST? AST это промежуточное представление программного кода в процессе компиляции, которое получается на выходе из парсера.

Например, этот код

              
def func(who):
    print "Hello, %s!" % who

func()

            

будет преобразован в следующее AST:

              
FunctionDef(
    name='func',  # имя функции
    args=arguments(  # дефолтные аргументы
        args=[Name(id='who', ctx=Param())],
        vararg=None,
        kwarg=None,
        defaults=[]),
    body=[  # тело функции
        Print(dest=None,
              values=[
                BinOp(left=Str(s='Hello %s!'),
                      op=Mod(),
                      right=Name(id='who', ctx=Load()))],
              nl=True)],
    decorator_list=[]),   # декораторы
Expr(value=Call(  # вызов функции
        func=Name(id='func', ctx=Load()),  # имя функции
        args=[],  # позиционные аргументы
        keywords=[],  # k-v аргументы
        starargs=None,  # *args аргументы
        kwargs=None))  # **kwargs аргументы

            

На первый взгляд ничего не понятно, но если приглядеться - то можно угадать назначение любого элемента этого дерева. AST сам по себе является Python объектом, так что им можно манипулировать как с любыми другими объектами. Полная документация по элементам и инструментам для работы с AST (имеются в стандартной библиотеке в модуле ast) есть тут.

Так вот, вернёмся к Tornado. Попробуем использовать такие-же обозначения как в оригинальной статье, т.е. декоратор с именем @shortgen и оператор бинарного сдвига <<.
Будем использовать тот же пример кода, что и в оригинальной статье.

Подготовка

Установим tornado

              
mkdir tornado-shortgen
cd tornado-shortgen/
virtualenv .env
source .env/bin/activate
pip install tornado

            

Напишем Tornado - приложение

              
import tornado.ioloop
import tornado.web
import tornado.gen
import os

class Handler(web.RequestHandler):
    @asynchronous
    @gen.engine
    @shortgen
    def get_short(self):
        (result, status) << self.db.posts.find_e({'name': 'post'})

    @asynchronous
    @gen.engine
    def get(self):
        (result, status) = yield gen.Task(self.db.posts.find_e, {'name': 'post'})

application = tornado.web.Application([
        (r"/", Handler),
])

if __name__ == "__main__":
    application.listen(8888)
    tornado.ioloop.IOLoop.instance().start()

            

Сохраним в файл shortgen_test.py

Реализация трансформации

Попробуем получить AST нашего модуля.

              
$ python
>>> import ast
>>> print ast.dump(ast.parse(open("shortgen_test.py").read()))

            

Увидим длинную неотформатированную портянку текста, из которой нас интересует только определения функций get_short и get

get_short - исходная функция с бинарным сдвигом и декоратором

              
FunctionDef(
    name='get_short',
    args=arguments(args=[Name(id='self', ctx=Param())],
                   vararg=None,
                   kwarg=None,
                   defaults=[]),
    body=[
        Expr(value=BinOp(  # операция с 2-мя операндами
                left=Tuple(  # левый операнд - это кортеж
                    elts=[Name(id='result', ctx=Load()), Name(id='status', ctx=Load())],
                    ctx=Load()),
                op=LShift(),  # операция бинарного сдвига
                right=Call(  # правый операнд - вызов функции self.db.posts.find_e
                    func=Attribute(
                        value=Attribute(
                            value=Attribute(
                                value=Name(id='self', ctx=Load()),
                                attr='db',
                                ctx=Load()),
                            attr='posts',
                            ctx=Load()),
                        attr='find_e',
                        ctx=Load()),
                    args=[Dict(keys=[Str(s='name')], values=[Str(s='post')])],  # функция вызывается с одним позиционным аргументом
                    keywords=[],
                    starargs=None,
                    kwargs=None)))],
    decorator_list=[  # список декораторов
        Attribute(value=Name(id='web', ctx=Load()), attr='asynchronous', ctx=Load()),
        Attribute(value=Name(id='gen', ctx=Load()), attr='engine', ctx=Load()),
        Name(id='shortgen', ctx=Load())])  # а вот и наш декоратор!

            

get - желаемый результат

              
FunctionDef(
    name='get',
    args=arguments(args=[Name(id='self', ctx=Param())],
                   vararg=None, kwarg=None, defaults=[]),
    body=[
        Assign(  # операция присваивания
            targets=[
                Tuple(elts=[ # с левой стороны от = находится такой-же tuple, но ctx изменился на Store()
                        Name(id='result', ctx=Store()), 
                        Name(id='status', ctx=Store())],
                      ctx=Store())],
            value=Yield(  # с правой - yield и вызов функции
                value=Call(  # вызов gen.Task
                    func=Attribute(
                        value=Name(id='gen', ctx=Load()),
                        attr='Task', ctx=Load()),
                    args=[Attribute(  # первый аргумент - имя функции self.db.posts.find_e
                            value=Attribute(
                                value=Attribute(
                                    value=Name(id='self', ctx=Load()),
                                    attr='db', ctx=Load()),
                                attr='posts', ctx=Load()),
                            attr='find_e', ctx=Load()),
                          Dict(keys=[Str(s='name')], values=[Str(s='post')])],
                    keywords=[],  # остальные аргументы не изменились
                    starargs=None,
                    kwargs=None)))],
    decorator_list=[
        Name(id='asynchronous', ctx=Load()),
        Attribute(value=Name(id='gen', ctx=Load()), attr='engine', ctx=Load())])  # декоратора shortgen нет

            

Выглядит монструозно, но зато как гибко! На самом деле всё просто.
Давайте посмотрим на различия:

  1. Полностью пропал Expr
  2. Вместо BinOp(left, op, right) теперь Assign(targets, value)
  3. У правого операнда значения ctx изменилось с Load на Store
  4. Вызов self.db.posts.find_e(...) заменен на gen.Task(self.db.posts.find_e, ...)
  5. Добавился Yield вокруг вызова функции
  6. Пропал декоратор @shortgen

Соответственно, чтобы получить из первого второе, нам нужно:

  1. Найти функцию, у которой в decorator_list есть декоратор @shortgen
  2. Удалить этот декоратор
  3. Найти в теле функции оператор бинарного сдвига BinOp
  4. Сохранить левый и правый операнды. В левом заменить ctx с Load на Store, из правого операнда извлечь название функции и её аргументы (позиционные, kw, и "звёздочные" - *, **)
  5. Добавить название функции ( self.db.posts.find_e) первым позиционным аргументом (т.е. в нашем примере получим позиционные аргументы [self.db.posts.find_e, {'name': 'post'}], а все остальные пустые
  6. Создать новый Call, но уже функции gen.Task с этими аргументами
  7. Обернуть его в Yield
  8. Создать Assign(targets, value) и в качестве targets взять сохраненный ранее левый операнд BinOp а в качестве value - только что созданный нами Yield
  9. Заменить в исходном дереве Expr на наш свежесобранный Assign

Хоть звучит сложно, но в коде это заняло чуть больше 50 строк. Если что-то не понятно - смотрите сразу туда.

Как это реализовать? Можно написать решение в лоб каким-то while циклом или рекурсивной функцией. Но мы воспользуемся паттерном Visitor и его адаптацией ast.NodeTransformer

Это класс, от которого можно отнаследоваться и насоздавать в нём методов типа visit_[NodeType] например visit_FunctionDef или visit_Expr. Значение, которое вернет метод, станет новым значением элемента AST. А сам Visitor просто рекурсивно обходит дерево, вызывая наши методы тогда, когда в дереве встретился соответствующий элемент. Это поможет нам логичнее организовать наш код.

  1. Создаем метод visit_FunctionDef, для отлова декорированной функции. В нём проверяем, что функция обернута в декоратор. Если обернута - удаляем декоратор и ставим пометку self.decorated
  2. Создаем метод visit_Expression, для отлова бинарного сдвига. В нем проверяем, что выставлен флаг self.decorated и что Expr - это именно бинарный сдвиг. Проводим остальные манипуляции (преобразование Expr в Assign) вручную. Благо, все нужные данные уже рядышком.

Собственно код на gist.github.com

Обратите внимание - декоратор shortgen при попытке его применения сразу выбрасывает Exception. Но если трансформация применилась успешно, то декоратор просто удаляется. Это защищает нас от случайных ошибок.

Полученный AST можно либо исполнить:

              
with open(filepath) as src:
    orig_ast = ast.parse(src.read())
new_ast = RewriteGenTask().visit(orig_ast)
code = compile(new_ast, filename, 'exec')
exec code

            

Либо сохранить в .pyo файл
http://stackoverflow.com/questions/8627835/generate-pyc-from-python-ast
https://gist.github.com/3849217#L172
и затем импортировать / вызывать python my_module.pyo

Заключение

Трансформация AST - более надежный и портируемый способ трансформации кода программы. Писать такие трансформации гораздо проще, чем модифицировать байткод. Этот способ широко применяется во многих языках, например в Lisp или Erlang.
Второй плюс - нет необходимости ничего манкипатчить, трансформация работает и с нашим и с внешним кодом одинаково.
Остальные плюсы и минусы расписаны в моём комментарии к оригинальной статье. Еще раз отмечу, что основной недостаток - проблематично применить трансформацию AST на лету. Она должна осуществляться на стадии компиляции в .pyc файл. (Ну и, конечно, если применяешь такие хаки, нужно это хорошо задокументировать).
Для маленьких проектов, в которых этот yield пишется в паре мест, такой сахар не имеет особого смысла, плюс усложняет разработку, т.к. появляется отдельный этап компиляции файла. Но на больших Tornado проектах можно и попробовать.

Ссылки

Весь код целиком на Gist
Документация по AST
Документация по tornado.gen
Генерация .pyc файла из AST
Если всё это кажется страшными костылями, есть выход xD

Домашнее задание

  • Согласитесь, этот список из 3-х декораторов выглядит жутковато?
    
    @asynchronous
    @gen.engine
    @shortgen
    

    Как с помощью AST обойтись всего одним @shortgen? А как без AST?

  • Текущая реализация требует, чтобы декоратор применялся именно как @shortgen ( @my_module.shortgen уже не сработает). Так же требуется, чтобы модуль tornado.gen был импортирован как from tornado import gen ( import tornado.gen или from tornado.gen import Task уже не прокатит). Как это исправить?
  • Попробуйте переписать сервер из статьи Раздача больших файлов через сервер TornadoWEB с использованием shortgen, скомпилировать и запустить.

Буду благодарен, если кто-то посоветует как такую трансформацию применять автоматически.