Самая коротка запись асинхронных вызовов в tornado или патчим байткод в декораторе

Сложный асинхронный обработчик в tornado иногда расползается на десятки callback функций, из-за чего становится трудно воспринимать и модифицировать код. Поэтому существует модуль tornado.gen, позволяющий писать обработчик как генератор. Но много yield gen.Task(...) тоже выглядит не очень. Поэтому в порыве бреда я написал упрощающий запись декоратор:

@asynchronous
@gen.engine
def get(self):
    result, status = yield gen.Task(
        db.users.find_one, {
            '_id': ObjectId(user_id),
        },
    )

@asynchronous
@gen.engine
@shortgen
def get(self):
    result, status << db.users.find_one_e({
        '_id': ObjectId(user_id),
        },
    )

Как это работает

Как вы уже заметили, мы заменили yield на <<. Так как python нам не позволит сделать это стандартными средствами, нам нужно модифицировать байткод. Для простой работы с ним воспользуемся модулем Byteplay. Посмотрим байткод двух простых функций:

from byteplay import Code
from pprint import pprint

def gen():
    a = yield 1
pprint(Code.from_code(gen.func_code).code)

[(SetLineno, 5),  # переходим на 5 строку
 (LOAD_CONST, 1), # загружаем константу 1
 (YIELD_VALUE, None), # "отдаём" загруженное значение
 (STORE_FAST, 'a'), # записываем в переменную a
 (LOAD_CONST, None),
 (RETURN_VALUE, None)]

def shift():
    a << 1
pprint(Code.from_code(shift.func_code).code)

[(SetLineno, 10),
 (LOAD_GLOBAL, 'a'),  # a из глобального пространства
 (LOAD_CONST, 1), # загружаем константу 1
 (BINARY_LSHIFT, None), # делаем сдвиг влево для a
 (POP_TOP, None),  # убираем верхний элемент стека
 (LOAD_CONST, None),
 (RETURN_VALUE, None)]

Поэтому сделаем простой патчер сугубо для этой ситуации:

from byteplay import YIELD_VALUE, STORE_FAST
code = Code.from_code(shift.func_code)
code.code[3] = (YIELD_VALUE, None)
code.code[4] = (STORE_FAST, 'a')
code.code.pop(1)
pprint(code.code)

[(SetLineno, 10), 
 (LOAD_CONST, 1),
 (YIELD_VALUE, None),
 (STORE_FAST, 'a'),
 (LOAD_CONST, None),
 (RETURN_VALUE, None)]

Теперь у нас есть байткод почти идентичный байткоду функции gen, применим его к shift и проверим результат:

shift.func_code = code.to_code()
res_gen = gen().send(None)
res_shift = shift().send(None)
print res_gen
print res_shift
print res_gen == res_shift

1
1
True

Результат получился одинаковым. Код для общей ситуации можно посмотреть на github. Про байткод подробнее можно узнать в официальной документации. А пока мы вернёмся к tornado. Возьмём уже готовый декоратор shortgen. И напишем простой обработчик:

def fetch(callback):
    callback(1)

class Handler(BaseHandler):
    @asynchronous
    @gen.engine
    @shortgen
    def get(self):
        result << gen.Task(fetch)

Код стал немного лучше, но нам всё равно приходится вручную оборачивать вызов в gen.Task, поэтому создадим ещё один декоратор для автоматизации этого процесса:

def fastgen(fnc):
    return partial(gen.Task, fnc)

@fastgen
def fetch(callback):
    callback(1)

class Handler(BaseHandler):
    @asynchronous
    @gen.engine
    @shortgen
    def get(self):
        result << fetch()

Теперь всё выглядит вполне прилично, но как это будет работать со сторонними библиотеками? А никак, поэтому теперь нам нужно пропатчить их! Нет, патчить байткод мы сейчас не будем, а применим просто monkey patch. Что бы не сломать старый код, мы заменим __getattribute__ у нужных классов на:

def getattribute(self, name):
    attr = None
    if name.find('_e') == len(name) - 2:
        attr = getattr(self, name[:-2])
    if hasattr(attr, '__call__'):
        return fastgen(attr)
    else:
        return super(self.__class__, self).__getattribute__(name)

Теперь если у пропатченного объекта нет атрибута, например, find_e(постфикс _e добавлен что бы не сломать старый код) нам вернётся атрибут find, обёрнутый в декоратор fasttgen.
И теперь код, например для asyncmongo, будет выглядеть так:

from asyncmongo.cursor import Cursor
Cursor.__getattribute__ = getattribute
class Handler(BaseHandler):
    @asynchronous
    @gen.engine
    @shortgen
    def get(self):
        result, status << self.db.posts.find_e({'name': 'post'})

Как этим воспользоваться

Для начала установим получившийся модуль:

pip install -e git+https://github.com/nvbn/evilshortgen.git#egg=evilshortgen

Теперь пропатчим нужные нам классы:

from evilshortgen import shortpatch
shortpatch(Cls1, Cls2, Cls3)

Обернём в декоратор собственные асинхронные методы и функции:

from evilshortgen import fastgen
@fastgen
def fetch(id, callback):
    return callback(id)

И воспользуемся в обработчике:

from evilshortgen import shortgen
class Handler(BaseHandler):
    @asynchronous
    @gen.engine
    @shortgen
    def get(self, id):
        data << fetch(12)
        num, user << Cls

Известные проблемы

Вызов может устанавливать значение только переменным:

a << fetch()  # работает
self.a << fetch()  # не работает

Сложные распаковки не поддерживаются:

a, b << fetch()  # работает
(a, b), c << fetch()  # не работает

Ссылки

Evilshortgen на github
Подробно про байткод
Byteplay

Автор: nvbn