Самая короткая запись асинхронных вызовов в tornado v2, или патчим 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 (имеются в стандартной библиотеке в модуле 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) вручную. Благо, все нужные данные уже рядышком.


    Собственно код
    # -*- coding: utf-8 -*-
    '''
    Created on 2012-10-07
    
    @author: Sergey <me@seriyps.ru>
    
    Альтернативный вариант решения из статьи http://habrahabr.ru/post/153595/
    на базе модификации AST
    '''
    import ast
    import marshal
    import py_compile
    import time
    import os.path
    
    
    class RewriteGenTask(ast.NodeTransformer):
    
        def __init__(self, *args, **kwargs):
            self.on_decorator = []
            self.on_assign = []
            super(RewriteGenTask, self).__init__(*args, **kwargs)
    
        def shortgen_deco_pos(self, decorator_list):
            # проверяет, что в списке декораторов имеется декоратор с именем
            # shortgen и возвращает его позицию.
            for pos, deco in enumerate(decorator_list):
                # Name(id='shortgen', ctx=Load())
                if isinstance(deco, ast.Name) and deco.id == 'shortgen':
                    return pos
            return -1
    
        def visit_FunctionDef(self, node):
            """
            Проверяет, что функция обернута в декоратор shortgen.
            Если обернута, удаляем декоратор и трансформируем содержимое.
            FunctionDef(
                name='get_short',
                args=arguments(...),
                body=[...],
                decorator_list=[
                    Attribute(value=Name(id='web', ...), attr='asynchronous', ...),
                    Attribute(value=Name(id='gen', ...), attr='engine', ...),
                    Name(id='shortgen', ctx=Load())])
            """
            deco_pos = self.shortgen_deco_pos(node.decorator_list)
            if deco_pos >= 0:
                # если функция обернута в shortgen декоратор, удаляем его,
                # делаем пометку в стеке и запускаем Visitor по содержимому
                # функции
                self.on_decorator.append(True)
                node.decorator_list.pop(deco_pos)
                self.generic_visit(node)  # трансформируем содержимое функции
                self.on_decorator.pop()
            return node
    
        def visit_Expr(self, expr):
            """
            == Основная трансформация ==
            Трансформируем
              result2 << func(arg, k=v, *args, **kwargs)
            в
              result2 = gen.Task(func, arg, k=v, *args, **kwargs)
    
            Пример AST представления "stmt << func(...)" (исходные данные):
            Expr(value=BinOp(left=Name(id='result', ctx=Load()),
                  op=LShift(),
                  right=Call(
                      func=Name(id='fetch', ctx=Load()),
                      args=[Num(n=1)],
                      keywords=[keyword(arg='k', value=Num(n=2))],
                      starargs=Tuple(elts=[Num(n=3)], ctx=Load()),
                      kwargs=Dict(keys=[Str(s='k2')], values=[Num(n=4)])))))
            ---- vvvvvvvvvvv ----
            Пример AST представления "stmt = yield func(...)" (результат):
            Assign(targets=[Name(id='result', ctx=Store())],
                   value=Yield(value=Call(
                       func=Attribute(value=Name(id='gen', ctx=Load()),
                                      attr='Task', ctx=Load()),
                       args=[Name(id='fetch', ctx=Load()), Num(n=1)],
                       keywords=[keyword(arg='k', value=Num(n=2))],
                       starargs=Tuple(elts=[Num(n=3)], ctx=Load()),
                       kwargs=Dict(keys=[Str(s='k2')], values=[Num(n=4)]))))
            """
            node = expr.value       # BinOp
            if not (self.on_decorator
                    and isinstance(expr.value, ast.BinOp)
                    and isinstance(node.op, ast.LShift)):
                # если функция не обернута в декоратор (on_decorator пуст), ничего
                # не меняем
                return expr
            # если функция, содержащая LShift, обернута в декоратор,
            # то заменяем на вызов gen.Task()
    
            # для начала конвертируем изменение на месте (stmt <<) на
            # присваивание (stmt =). Для этого заменяем ctx=Load на
            # ctx=Store (см self.visit_Load())
            self.on_assign.append(True)
            assign_target = self.visit(node.left)
            self.on_assign.pop()
            # генерируем присваивание ... = ...
            (new_node, ) = ast.Assign(
                targets = [assign_target],
                value = ast.Yield(
                    value=self.construct_gen_task_call(node.right))),
            # копируем номер линии оригинальной конструкции
            new_node = ast.fix_missing_locations(ast.copy_location(new_node, expr))
            return new_node
    
        def construct_gen_task_call(self, func_call):
            """
            Конвертируем вызов функции в вызов gen.Task с именем функции первым
            параметром
              func(arg, k=v, *args, **kwargs)
            в
              gen.Task(func, arg, k=v, *args, **kwargs)
            Пример AST представления "func(...)":
            Call(
                func=Name(id='fetch', ctx=Load()),
                args=[Num(n=1)],
                keywords=[keyword(arg='k', value=Num(n=2))],
                starargs=Tuple(elts=[Num(n=3)], ctx=Load()),
                kwargs=Dict(keys=[Str(s='k2')], values=[Num(n=4)])))
            ---- vvvvvvvvv ----
            Пример AST представления "gen.Task(func, ...)":
            Call(
                func=Attribute(value=Name(id='gen', ctx=Load()),
                               attr='Task', ctx=Load()),
                args=[Name(id='fetch', ctx=Load()), Num(n=1)],
                keywords=[keyword(arg='k', value=Num(n=2))],
                starargs=Tuple(elts=[Num(n=3)], ctx=Load()),
                kwargs=Dict(keys=[Str(s='k2')], values=[Num(n=4)]))
            """
            # Генерируем gen.Task
            gen_task = ast.Attribute(
                    value=ast.Name(id='gen', ctx=ast.Load()),
                    attr='Task', ctx=ast.Load())
            # Генерируем вызов gen.Task(func, ...)
            call = ast.Call(
                func=gen_task,
                # имя оригинальной ф-ции 1-м аргументом:
                args=[func_call.func] + func_call.args,
                keywords=func_call.keywords,
                starargs=func_call.starargs,
                kwargs=func_call.kwargs)
            return self.visit(call)
    
        def visit_Load(self, node):
            # Заменяем Load() на Store()
            if self.on_assign:
                return ast.copy_location(ast.Store(), node)
            return node
    
    
    def shortgen(f):
        raise RuntimeError("ERROR! file must be compiled with yield_ast!")
    
    
    def compile_file(filepath):
        path, filename = os.path.split(filepath)
        with open(filepath) as src:
            orig_ast = ast.parse(src.read())
        new_ast = RewriteGenTask().visit(orig_ast)
        code = compile(new_ast, filename, 'exec')
    
        pyc_filename = os.path.splitext(filename)[0] + '.pyc'
        pyc_filepath = os.path.join(path, pyc_filename)
    
        with open(pyc_filepath, 'wb') as fc:
            fc.write(py_compile.MAGIC)
            py_compile.wr_long(fc, long(time.time()))
            marshal.dump(code, fc)
            fc.flush()
    
    if __name__ == '__main__':
        import sys
        if len(sys.argv) < 2:
            print "Usage: %s file_to_compile1.py [file2.py] ..." % sys.argv[0]
        for filename in sys.argv[1:]:
            compile_file(filename)
    
    

    Gist

    Полученный 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 файл
    stackoverflow.com/questions/8627835/generate-pyc-from-python-ast
    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, скомпилировать и запустить.
    • +31
    • 3,9k
    • 4
    Поделиться публикацией
    Реклама помогает поддерживать и развивать наши сервисы

    Подробнее
    Реклама
    Комментарии 4
    • НЛО прилетело и опубликовало эту надпись здесь
      • +2
        Сам испугался — как такую простую вещь умудрился так сложно описать.
        • +1
          Та не, круто, например в новом C# для этого пришлось новые ключевые слова вводить async/await.
      • 0
        Да, такой вариант менее костыльно выглядит =)

        Только полноправные пользователи могут оставлять комментарии. Войдите, пожалуйста.