Как я переизобрел словари в Python

В нашем Django-приложении необходимо было разработать отчет (расчет) бонусов.
Отчет должен иметь вложенную структуру с подведением итогов по пользователям, подразделениям и по всей компании. Схематично его логику можно представить:

print total
for department in departments:
    print department.total
    for user in department.users:
        print user.total
        for row in user.rows:
            print row.data

У этого отчета было два осложняющих момента:

1. В роли "row" могли выступать разные модели (и располагаться вперемежку), что не позволяет использовать итераторы по QuerySet'ам.
2. Время построение отчета. Сбор данных занимает существенное время (несколько секунд). Данные в отчете могут меняться. Говоря на чистоту, это не статический отчет, а инструмент для контроля и корректировки начисленных бонусов в виде отчета. Но данные меняются не очень часто, скажем на каждые 100 просмотров придется одно изменение, после которого нужно перестроить отчет. Т.е. данные можно кэшировать.

Структура из вложенных словарей отлично решает обе задачи: в них можно сложить все требуемые скаляры (числа, строки, даты), сериализовать и сложить в кэш.

Структура данных для отчета приобрела вид (упрощена):

{
    'total': {
        'income': 1234,
        'bonus': 123,
        'expense': 1234,
        'penalty': 123
    },
    'departments': {
        '{dept_id}': {
            'department': {
                'title': 'Mega Department'
            }
            'total': {
                'income': 1234,
                'bonus': 123,
                'expense': 1234,
                'penalty': 123
            },
            'users': {
                '{user_id}': {
                    'user': {
                        'name': 'John Smith'
                    },
                    'total': {
                        'income': 1234,
                        'bonus': 123,
                        'expense': 1234,
                        'penalty': 123
                    },
                    'rows': {
                        '{sale_id}': {        //  Одна модель
                            'type': 'sale'
                            'base_income': 1234,
                            'bonus': 123,
                            'comment': 'some description'
                        },
                        '{expense_id}': {     //  Другая модель !!!
                            'type': 'expense'
                            'expense': 1234,
                            'penalty': 123,
                            'comment': 'some description'
                        },
                        ...
                    }
                },
                ...
            }
        },
        ...
    }
}

И вот тут-то я столкнулся с проблемой, что заполнение такой структуры из словарей не столь удобно, как мне того хотелось. Проверка словарей на наличие ключей или использование setdefatult(key, {}) превращает код в нечитабельную кашу.

Эта структура чем-то напоминает XML. И мне бы хотелось использовать что-то подобное тому, как строятся XPath-выражения для адресации узлов XML-дерева:

/departments/{dept_id}/users/{user_id}/rows/{row_id}/base_income

или на языке Python что-то вида:

data.departments.{dept_id}.users.{user_id}.rows.{row_id}.base_income

Учтывая, что {dept_id} и прочие другие {id} — целые числа, то я разрешил себе использование квадратных скобок: [].

data.departments[{dept_id}].users[{user_id}].rows[{row_id}].base_income

Собственно мне нужен был такой класс, который бы вел себя в основном, как словарь, но при этом:

1. доступ к атрибутам можно было делать без квадратных скобочек
2. автоматически создавались отсутствующие аттрибуты

Так появился ElasticDict

В итоге

Код по подготовке данных выглядит приблизительно так:

data = ElasticDict()
for sale in Sale.objects.filter(...).prefetch_related(...):
    data.departments[sale.user.department.pk].users[sale.user.pk].rows[sale.pk] = {'base_income': sale.amount, 'bonus': sale.calc_bonus()}

# или в другой форме, кому как больше нравится
for expense in Expense.objects.filter(...).prefetch_related(...):
    data.departments[sale.user.department.pk].users[sale.user.pk].rows[expense.pk].base_expense = expense.amount
    data.departments[sale.user.department.pk].users[sale.user.pk].rows[expense.pk].penalty = expense.calc_penalty()

Код в шаблоне так:

{{ data.total }}
{% for dept_id, department in data.departments.items %}
    {{ department.total }}
    {% for user_id, user in department.users.items %}
        {{ user.total }}
        {% for row_id, row in user.rows.items %}:
            {{ row.data }}
        {% endfor %}
    {% endfor %}
{% endfor %}

Заключение

Надо отметить, что ElasticDict() это подкласс обычного dict()'а, т.е. в нем доступно все то, что и в обычном словаре. В тот момент, когда потребуется "зафиксировать" структуру (снова захотим получать KeyError'ы при обращении к несуществующим ключам), экземпляр ElasticDict можно экспортировать в обчный dict(). Делается рекурсивный обход ElasticDict(), где все экземпляры этого класа заменяются на обычные словари. Есть и обратное преобразование — на вход подаем словарь, на выходе получаем ElasticDict также с рекурсивным обходом.

Замечания/предложения приветствуются!

Автор: Ruzin

Источник