Генераторы и сопрограммы Python. Часть 2

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

На сей раз мы ещё немного усложним задачу.

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

Во-вторых, теперь перед добавлением записи нам нужно проверить, существует ли она. Если запись уже существует, то её не нужно добавлять.

Для решения первой задачи сначала мне пришло в голову такое решение - сделать дополнительный класс-обёртку с методом __call__:

class wrapped_coro():
    def __init__(self, coro):
        self.coro = coro

    def __call__(self, *args, **kwargs):
        return self.coro.send(*args, **kwargs)

class wrapper():
    def __init__(self, coro, *args, **kwargs):
        self.coro = coro(*args, **kwargs)

    def __enter__(self):
        self.coro.next()
        return wrapped_coro(self.coro)

    def __exit__(self, type, value, traceback):
        self.coro.close()
        if value is None:
            return True
        return False

Однако, немного подумав, я понял, что это ничем не оправданный оверинжениринг, удалил класс wrapped_coro и переписал wrapper вот так:

class wrapper():
    def __init__(self, coro, *args, **kwargs):
        self.coro = coro(*args, **kwargs)

    def __enter__(self):
        self.coro.next()
        return self.coro.send

    def __exit__(self, type, value, traceback):
        self.coro.close()
        if value is None:
            return True
        return False

Получилось даже ещё проще, чем в прошлой заметке. Теперь функцию копирования можно переписать так:

def copy(db):
    """
    Подпрогрмма, использующая генератор и сопрограмму для копирования
    содержимого таблицы user в таблицу user2
    """
    with wrapper(writer, db) as write:
        for row in reader(db):
            write(row)

Для решения второй задачи воспользуемся ещё одной возможностью, предоставляемой оператором yield - он может не только возвращать значение или только считывать его, но и считывать и возвращать одновременно. Делается это так:

def checker(db):
    """Сопрограмма. Проверяет, что указанный пользователь уже добавлен в таблицу users2"""
    select = db.cursor()
    exists = None
    try:
        while True:
            row = (yield exists)
            select.execute('''SELECT COUNT(*)
                              FROM user2
                              WHERE surname = %s
                                AND name = %s
                                AND patronym = %s''', row)
            count, = select.fetchone()
            exists = count != 0
    except GeneratorExit:
        select.close()

Перед началом использованием сопрограммы, как и прежде, нужно прокрутить её до первого оператора yield. Делается это, как и прежде, вызовом метода next из обёртки.

Новый вариант функции копирования примет следующий вид:

def copy(db):
    """
    Подпрогрмма, использующая сопрограммы для дополнения таблицы user2
    содержимым таблицы user
    """
    with wrapper(writer, db) as write:
        with wrapper(checker, db) as check:
            for row in reader(db):
                if not check(row):
                    write(row)

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

#!/usr/bin/python
# -*- coding: UTF-8 -*-

import MySQLdb

def reader(db):
    """Генератор. Читает строки таблицы user"""
    select = db.cursor()
    select.execute('SELECT surname, name, patronym FROM user')
    for row in select:
        yield row
    select.close()

def writer(db):
    """Сопрограмма. Пишет строки в таблицу user2"""
    insert = db.cursor()
    try:
        while True:
            row = (yield)
            try:
                insert.execute('INSERT INTO user2(surname, name, patronym) VALUES(%s, %s, %s)', row)
                db.commit()
            except:
                db.rollback()
    except GeneratorExit:
        insert.close()

def checker(db):
    """Сопрограмма. Проверяет, что указанный пользователь уже добавлен в таблицу users2"""
    select = db.cursor()
    exists = None
    try:
        while True:
            row = (yield exists)
            select.execute('''SELECT COUNT(*)
                              FROM user2
                              WHERE surname = %s
                                AND name = %s
                                AND patronym = %s''', row)
            count, = select.fetchone()
            exists = count != 0
    except GeneratorExit:
        select.close()

class wrapper():
    def __init__(self, coro, *args, **kwargs):
        self.coro = coro(*args, **kwargs)

    def __enter__(self):
        self.coro.next()
        return self.coro.send

    def __exit__(self, type, value, traceback):
        self.coro.close()
        if value is None:
            return True
        return False

def copy(db):
    """
    Подпрогрмма, использующая сопрограммы для дополнения таблицы user2
    содержимым таблицы user
    """
    with wrapper(writer, db) as write:
        with wrapper(checker, db) as check:
            for row in reader(db):
                if not check(row):
                    write(row)

db = MySQLdb.connect(user = 'user',
                     passwd = 'p4ssw0rd',
                     db = 'database',
                     charset = 'UTF8')

copy(db)

db.close()