Использование typing.Generic в Python

Я работаю над проектом с довольно большой кодовой базой. Проект с историей.

Некоторые части наша команда написала задолго до аннотаций типов. Мы до сих пор добавляем их в наш легаси‑код и улучшаем существующие подсказки. Стоит ли эта игра свеч? Определённо — да.

Наши пользователи — разработчики. Они открывают наш код в PyCharm ежедневно и надеются, что он поможет им решить их задачи максимально быстро и просто.

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

Один из моих коллег добавил аннотацию типов, которая выглядит так:

# file name: type_cast_0.py
class A:
    a = 'a'

class B(A):
    b = 'b'

class DoSomethingWithA:
    _class = A
    def do(self) -> A:
        return self._class()

class DoSomethingWithB(DoSomethingWithA):
    _class = B

PyCharm не видит проблем в этом коде. Его анализатор показывает зелёную галочку. Mypy также не находит никаких проблем:

$ mypy type_cast_0.py
Success: no issues found in 1 source file

Но если добавить вот такой код, который использует DoSomethingWithB…

# file name: type_cast_1.py
from type_cast_0 import DoSomethingWithB

print(DoSomethingWithB().do().b)

PyCharm теперь показывает warning: Unresolved attribute reference 'b' for class 'A'. И Mypy помечает этот кусок кода ошибкой:

$ mypy type_cast_1.py
type_cast_1.py:4: error: "A" has no attribute "b"
Found 1 error in 1 file (checked 1 source file)

Попробуем это исправить. Ниже — моя первая попытка: наивный подход к дженерикам в Python.

# file name: type_cast_2.py
#...
TV = tp.TypeVar('TV')

class DoSomethingWithA(tp.Generic[TV]):
    _class: tp.Type[TV] = A
    def do(self) -> TV:
        return self._class()

class DoSomethingWithB(DoSomethingWithA):
    _class = B

PyCharm не показывает никаких ошибок или предупреждений. Mypy всё ещё не нравится мой код:

$ mypy type_cast_3.py
type_cast_2.py:17: error: Incompatible types in assignment (expression has type "Type[A]", variable has type "Type[TV]")
Found 1 error in 1 file (checked 1 source file)

Интересно… Попробуем поменять TV = tp.TypeVar('TV') на TV = tp.TypeVar('TV', bound=A). Такая же ошибка. Становится интереснее…

Официальная документация не сильно помогает. В ней всего пара примеров использования Generics, но ничего, что даст ключ к исправлению проблемы. К счастью, есть прекрасный раздел о Generics в документации mypy.

Для моего примера код может выглядеть как‑то так:

# file name: type_cast_6.py
# ...
class DoSomethingWith(tp.Generic[TV]):
    _class: tp.Type[TV]

    def do(self) -> TV:
        return self._class()

А вот пример его использования:

# file name: type_cast_7.py
from type_cast_6 import DoSomethingWith, B

print(DoSomethingWith[B]().do().b)

Mypy не видит никаких проблем. PyCharm показывает зелёную галочку:

$ mypy type_cast_6.py
Success: no issues found in 1 source file
$ mypy type_cast_7.py
Success: no issues found in 1 source file

К сожалению, попытка выполнить этот код завалится с исключением:

$ python type_cast_7.py
...
AttributeError: 'DoSomethingWith' object has no attribute '_class'

В Python нет возможности использовать TypeVar так же, как можно использовать дженерики в Java, например. Я не могу присвоить TV переменной _class и ожидать, что Python заменит переменную типа на реальный класс во время выполнения. Другими словами, если использовать _class: tp.Type[TV] = TV в type_cast_6.py, я получу TypeError: 'TypeVar' object is not callable.

Чтобы этого избежать, я добавил подклассы для DoSomethingWith:

# file name: type_cast_8.py
# ...
class DoSomethingWithA(DoSomethingWith):
    _class = A

class DoSomethingWithB(DoSomethingWith):
    _class = B

# file name: type_cast_9.py
from type_cast_8 import DoSomethingWithB

print(DoSomethingWithB().do().b)

Не особенно элегантное решение, но оно работает.

В этом посте много примеров кода — я их порезал. Полные примеры можно найти в специальном репозитории на моём аккаунте в GitHub.