Типизируй все

Всем привет!

У нас уже есть одна статья про развитие типизации в Ostrovok.ru. В ней объясняется, зачем мы переходим с pyContracts на typeguard, почему переходим именно на typeguard и что в итоге получаем. А сегодня я расскажу подробнее о том, каким образом происходит этот переход.

Объявление функции с pyContracts в общем случае выглядит так:

from contracts import new_contract
import datetime

@new_contract
def User (x):
    from models import User
    return isinstance(x, User)

@new_contract
def dt_datetime (x):
    return isinstance(x, datetime.datetime)

@contract
def func(user_list, amount, dt=None):
    """
    :type user_list: list(User)
    :type amount: int|float
    :type dt: dt_datetime|None
    :rtype: bool
    """
    …

Это абстрактный пример, потому что я не нашла в нашем проекте определения функции, короткого и содержательного по количеству случаев для проверки типа. Обычно определения для pyContracts хранятся в файлах, не содержащих никакой другой логики. Обратите внимание, что здесь User – это определенный пользовательский класс, и он не импортируется напрямую.

А это желаемый результат с typeguard:

from typechecked import typechecked
from typing import List, Optional, Union
from models import User
import datetime

@typechecked
def func (user_list: List[User], amount: Union[int, float], dt: Optional[datetime.datetime]=None) -> bool:
    ...

Вообще функций и методов с проверкой типа в проекте так много, что если сложить их в стопку, то можно дотянуться до Луны. Так что переводить их с pyContracts на typeguard вручную просто невозможно (я пробовала!). Поэтому я решила написать скрипт.

Скрипт разделяется на две части: одна кэширует импорты новых контрактов, а вторая занимается рефакторингом кода.

Хочу отметить, что ни тот, ни другой скрипт не претендует на универсальность. Мы не ставили целью написание инструмента для решения всех требуемых случаев. Поэтому я часто опускала автоматическую обработку каких-то частных случаев, если они редко встречаются в проекте — быстрее поправить руками. К примеру, скрипт генерации маппинга контрактов и импортов собрал 90% значений, оставшиеся 10% — это крафтовые маппинги ручной работы.

Логика работы скрипта для генерации маппинга:

Шаг 1. Пройтись по всем файлам проекта, прочитать их. Для каждого файла:

• если подстроки "@new_contract" нет, пропустить этот файл,
• если есть, то разбить файл по строке "@new_contract". Для каждого элемента:
  – распарсить на определение и импорт,
  – если получилось, записать в файл успехов,
  – если нет, записать в файл ошибок.

Шаг 2. Вручную обработать ошибки

Теперь, когда у нас есть имена всех типов, которыми пользуется pyContracts (это они и были определены с декоратором new_contract), и есть все необходимые импорты, можно писать код для рефакторинга. Пока я переводила с pyContracts на typeguard вручную, я поняла, что мне нужно от скрипта:

1. Это команда, которая принимает аргументом имя модуля (можно несколько), в котором надо заменить синтаксис аннотаций функций.
2. Пройтись по всем файлам модуля, прочитать их. Для каждого файла:
   
   • если подстроки “@contract” нет, пропустить этот файл;
   • если есть, то превратить код в ast (абстрактное синтаксическое дерево);
   • найти все функции, которые находятся под декоратором contract, для каждой:
     
     • получить докстринг, распарсить, потом удалить,
     • создать словарь вида {arg_name: arg_type}, с его помощью заменить аннотацию функции,
     • запомнить новые импорты,
   • модифицированное дерево записать в файл через astunparse;
   • добавить новые импорты в начало файла;
   • заменить строки "@contract" на "@typechecked" потому что так проще, чем через ast.

Решать вопрос "а не импортируется ли уже данное имя в данном файле?" я изначально не собиралась: с этой проблемой мы справимся дополнительным прогоном библиотеки isort.

Зато после прогона первой версии скрипта возникли вопросы, которые решать все же пришлось. Оказалось, что 1) ast не всемогущ, 2) astunparse более всемогущ, чем хотелось бы. Это проявлялось в следующем:

• в момент перехода к синтаксическому дереву из кода пропадают все однострочные комментарии;
• пустые строки тоже пропадают;
• ast не различает функции и методы класса, пришлось добавить логику;
• обратно, при переходе от дерева к коду многострочные комментарии в тройных кавычках записываются комментариями в одинарных кавычках и занимают одну строку, а переносы на новую строку заменяются на \n;
• появляются ненужные скобки, например if A and B and C or D становится if ((A and B and C) or D).

Код, пропущенный через ast и astunparse, остается рабочим, но его читаемость снижается.

Самый серьезный недостаток из перечисленных — это исчезающие однострочные комментарии (в других случаях мы ничего не теряем, а только приобретаем — скобочки, например). С этим обещает разобраться библиотека horast, основанная на ast, astunparse и tokenize. Обещает и делает.

Теперь пустые строки. Вариантов решения было два:

1. tokenize умеет определять “часть речи” питона, и этим пользуется horast, когда достает токены типа comment. Но tokenize так же имеет токены типа NewLine и NL. Значит, надо посмотреть, как horast восстанавливает комментарии, и скопировать, заменив тип токена.
   — предложила Аня, опыт в разработке 2 месяца
2. Раз horast может восстанавливать комментарии, то сначала заменим все пустые строки на определенный комментарий, потом пропустим через horast, и заменим наш комментарий на пустую строку.
   — придумал Женя, опыт в разработке 8 годиков

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

В horast’е мы пользуемся двумя функциями: parse и unparse, но обе неидеальны — parse содержит странные внутренние ошибки, в редких случаях не может распарсить исходный код, а unparse не может записать что-то, что имеет тип type (такой тип, который получится, если сделать type(any_other_type)).

С parse я решила не разбираться, потому что логика работы довольно запутанная, а исключения случаются редко — здесь работает принцип неуниверсальности.

А вот unparse работает предельно ясно и довольно изящно. Функция unparse создает экземпляр класса Unparser, который в init обрабатывает дерево, а потом записывает его в файл. Horast.Unparser последовательно наследуется от многих других Unparser’ов, где самый базовый класс — это astunparse.Unparser. Все классы-наследники просто расширяют функционал базового класса, но логика работы остается такой же, так что рассмотрим astunparse.Unparser. В нем есть пять важных методов:

1. write – просто записывает что-то в файл.
2. fill – использует write с учетом количества отступов (количество отступов хранится как поле класса).
3. enter – увеличивает отступ.
4. leave – уменьшает отступ.
5. dispatch – определяет тип узла дерева (допустим, T), вызывает соответствующий ему метод по имени типа узла, но с нижним подчеркиванием (то есть _T). Это мета-метод.

Все остальные методы — это методы вида _T, например, _Module или _Str. В каждом таком методе может: 1) рекурсивно вызываться dispatch для узлов поддерева или 2) использоваться write для записи содержимого узла или добавления символов и ключевых слов, чтобы результат был валидным выражением на python.

Например, нам попался узел типа arg, в котором ast хранит имя аргумента и узел аннотации. Тогда dispatch вызовет метод _arg, который сначала запишет имя аргумента, потом запишет двоеточие и запустит dispatch для узла аннотации, где будет разбираться поддерево аннотации, и для этого поддерева все так же будут вызываться dispatch и write.

Вернемся к нашей проблеме невозможности обработки типа type. Теперь, когда понятно, как работает unparse, создать свой тип несложно. Создадим некий тип:

class NewType(object):
    def __init__ (self, t):
        self.s = t.s

Он хранит в себе строку, и не просто так: нам же нужно типизировать аргументы функций, а типы аргументов мы как раз получаем в виде строк из докстринга. Поэтому давайте подменим аннотации аргументов не теми типами, которые нам требуются, а объектом NewType, хранящем внутри только имя нужного типа.

Для этого расширим horast.Unparser – напишем свой UnparserWithType, отнаследовавшись от horast.Unparser, и добавим обработку нашего нового типа.

class UnparserWithType(horast.Unparser):
    def _NewType (self, t):
        self.write(t.s)

Это сочетается с духом библиотеки. Названия переменных выполнены в стилистике ast, и именно поэтому они состоят из одной буквы, а не потому что я не умею придумывать названия. Думаю, что t – это сокращение от tree, а s – от string. Кстати, NewType – это не строка. Если бы мы хотели, чтобы он интерпретировался как тип строки, то нам надо было бы до и после вызова write записать кавычки.

А теперь магия monkey patch: заменим horast.Unparser нашим UnparserWithType.

Как это теперь работает: у нас есть синтаксическое дерево, в нем есть какая-то функция, в функции — аргументы, у аргументов — аннотации типов, в аннотации типов спрятана игла, а в ней — смерть Кощеева. Раньше узлов аннотаций вообще не было, это мы их создали, причем любой такой узел есть экземпляр NewType. Мы вызываем функцию unparse для нашего дерева, и она для каждого узла вызывает dispatch, которая классифицирует этот узел и вызывает соответствующую ему функцию. Как только функция dispatch получает узел аргумента, она записывает имя аргумента, потом смотрит, есть ли аннотация (раньше это был None, но мы положили туда NewType), если есть, то пишет двоеточие и вызывает dispatch для аннотации, который вызывает наш _NewType, который просто записывает строку, которую хранит — это имя типа. В итоге получаем записанным аргумент: тип.

Вообще-то, это не совсем легально. С точки зрения компилятора, мы записали аннотации аргументов какими-то словами, которые нигде не определены, так что когда unparse завершает свою работу, мы получаем неправильный код: нам нужны импорты. Я просто формирую строку правильного формата и добавляю ее в начало файла, а потом дописываю результат unparse, хотя могла бы добавить импорты и как узлы в синтаксическое дерево, так как ast поддерживает узлы Import и ImportFrom.

Решение проблемы тройных кавычек не сложнее добавления нового типа. Мы создадим класс StrType и метод _StrType. Метод ничем не отличается от использовавшегося ранее для аннотации типов метода _NewType, а вот определение класса изменилось: мы будем хранить не только саму строку, но и уровень табуляции, на котором ее следует записывать. Число отступов определим так: если эта строка встретилась в функции, то один, если в методе, то два, а случаев, когда функция определена в теле другой функции и при этом декорирована, в нашем проекте нет.

class StrType(object):
    def __init__ (self, s, indent):
        self.s = s
        self.indent = indent

    def __repr__ (self):
         return '"""\n' + self.s + '\n' + ' ' * 4 * self.indent + '"""\n'

В repr определим, как должна выглядеть наша строка. Я думаю, это далеко не единственное решение, но оно работает. Можно было бы поэкспериментировать с astunparse.fill и astunparse.Unparser.indent, тогда это было бы более универсально, но эта идея пришла мне в голову уже во время написания этой статьи.

На этом решенные трудности заканчиваются. После работы моего скрипта иногда возникает проблема циклических импортов, но это вопрос архитектуры. Я не нашла готового стороннего решения, а обрабатывать такие случаи в рамках моего скрипта кажется серьезным усложнением поставленной задачи. Наверное, с помощью ast возможно обнаруживать и разрешать циклические импорты, но эту идею нужно обдумать отдельно. В общем, ничтожно малое количество таких казусов в нашем проекте вполне позволяло мне не обрабатывать их автоматически.

Еще одна сложность, которая мне встретилась — это отсутствие в ast обработки выражений from A import B as C. Внимательный читатель уже знает, что monkey patch — лекарство от всех болезней. Пусть это будет для него домашним заданием, а я решила сделать так: просто добавить такие импорты в файл маппинга, потому что обычно эта конструкция используется для обхода конфликта имен, а у нас их мало.

Несмотря на найденные несовершенства, скрипт делает то, ради чего задумывался. Что в итоге:

1. Время, за которое запускается проект, сократилось с 10 до 3 секунд;
2. Уменьшилось число файлов за счет удаления определений new_contract. Сократились сами файлы: я не замеряла, но в среднем гит насчитывал n добавленных строк и 2n удаленных;
3. Умные IDE стали делать разные подсказки, поскольку теперь это не комментарии, а честные импорты;
4. Улучшилась читаемость;
5. Кое-где появились скобочки.

Спасибо!

Полезные ссылки:

1. Ast
2. Horast
3. Все типы узлов ast и что в них хранится
4. Красиво показывает синтаксическое дерево
5. Isort