Оборачивание обобщенной функции в Python: почему теряется тип и что делать

Разбираем, почему при оборачивании обобщенной функции в Python тип T схлопывается в Never в mypy/pyright и как сохранить тип через дескриптор и декоратор.

Как обернуть обобщенную функцию в Python и не потерять ее тип: что идет не так и как с этим жить

Превратить функцию типа T → T в безаргументный thunk, который возвращает держатель для последующего выполнения, кажется простой задачей. Идея прозрачна: принять вызываемый объект, сохранить его в контейнере и вызвать, когда понадобится. Сложности начинаются, когда сам вызываемый объект — обобщенный. Средства статической проверки типов, такие как mypy и pyright, в этой схеме не сохраняют переменную типа, и итоговый тип схлопывается в Never. Разберемся, почему так происходит и что можно сделать уже сейчас.

Исходные данные

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

from collections.abc import Callable

class Hold[U]:
    def __init__(self, stored: Callable[[U], U]):
        self.stored = stored

def packer[U](fn: Callable[[U], U]) -> Callable[[], Hold[U]]:
    def make() -> Hold[U]:
        return Hold(fn)
    return make

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

@packer
def echo_str(arg: str) -> str:
    return arg

reveal_type(echo_str)
reveal_type(echo_str())
reveal_type(echo_str().stored)
reveal_type(echo_str().stored("test"))

Но при применении к обобщенной функции типовый аргумент пропадает, и вся цепочка выводится как Never.

@packer
def echo_generic[T](arg: T) -> T:
    return arg

reveal_type(echo_generic)
reveal_type(echo_generic())
reveal_type(echo_generic().stored)
reveal_type(echo_generic().stored("test"))

Логично ожидать, что echo_generic будет иметь тип def[T]() -> Hold[T], echo_generic().stored — def[T](T) -> T, а последний вызов разрешится в str. Однако ни mypy, ни pyright сейчас так не делают.

Что происходит на самом деле

Аннотации пытаются повторно использовать один и тот же символ T в разных областях видимости переменных типов. Обобщенные функции и обобщенные классы связывают переменные типов внутри своих тел. Переменные уровня модуля, однако, не могут нести свободную переменную типа. В данном случае после декорирования echo_generic становится переменной уровня модуля, и тип вроде Hold[T] на уровне модуля несостоятелен, потому что в этой области нет связывателя для T.

Поэтому проверяющий получает Hold[Never]. Это обозначает случай, который нельзя параметризовать конкретным типом на уровне модуля. Здесь напрямую проявляется разница между def (T) -> T для некоторого T и def[T](T) -> T для любого T. Эти формы принципиально различаются, и в системе типов Python нет единого интерфейса, который их объединяет. Это также вопрос раннего против позднего связывания. Когда применяется декоратор, его переменные типов должны быть разрешены, и текущие проверяющие не проводят обобщенные функции через этот этап. Спецификация подобного поведения не требует.

Прагматичный обходной путь с дескриптором

Если на уровне модуля доступен конкретный тип, сохраненный вызываемый объект обязан принимать и возвращать именно этот конкретный тип. Если конкретного типа нет, максимум, что можно сделать, — предоставить вызываемый объект, который возвращает то же, что получил. Это различие можно смоделировать дескриптором, меняющим тип атрибута stored в зависимости от того, параметризован ли Hold конкретным типом или нет.

Следующий шаблон кодирует это поведение. Он опирается на механизмы, активные только при TYPE_CHECKING, и перегрузки для наведения проверяющего, не меняя поведение во время выполнения. Учтите, это хак, зависящий от текущего вывода Hold[Never]. Если когда-нибудь проверяющий начнет выводить Hold[Any], он перестанет работать как задумано. Кроме того, он не распространяет границы или ограничения переменных типов; в обобщенной форме вызываемого они становятся неограниченными.

from collections.abc import Callable
from typing import TYPE_CHECKING, Any, Never, Self, overload


if TYPE_CHECKING:

    class StoredView:
        @overload  # type: ignore[no-overload-impl]
        def __get__(self, instance: None, owner: type[object], /) -> Self: ...
        @overload
        def __get__[R](
            self, instance: Hold[Never], owner: type[Hold[Never]], /
        ) -> Callable[[R], R]:
            """
            Конкретная параметризация недоступна; вернуть вызываемый объект,
            который просто возвращает тот же тип, что и получил
            """
        @overload
        def __get__[U](
            self, instance: Hold[U], owner: type[Hold[U]], /
        ) -> Callable[[U], U]:
            """
            Есть конкретная параметризация; вернуть вызываемый объект,
            ограниченный этим конкретным типом
            """
        def __set__[U](
            self, instance: Hold[Any], value: Callable[[U], U], /
        ) -> None: ...


class Hold[U]:
    if TYPE_CHECKING:
        stored = StoredView()

    def __init__(self, stored: Callable[[U], U]):
        self.stored = stored


def packer[U](fn: Callable[[U], U]) -> Callable[[], Hold[U]]:
    def make() -> Hold[U]:
        return Hold(fn)
    return make


@packer
def echo_str(arg: str) -> str:
    return arg

reveal_type(echo_str)
reveal_type(echo_str())
reveal_type(echo_str().stored)
reveal_type(echo_str().stored("test"))
reveal_type(echo_str().stored(1))


@packer
def echo_generic[T](arg: T) -> T:
    return arg

reveal_type(echo_generic)
reveal_type(echo_generic())
reveal_type(echo_generic().stored)
reveal_type(echo_generic().stored("test"))
reveal_type(echo_generic().stored(1))

С такой настройкой, когда конкретный тип известен, stored типизируется как вызываемый объект, принимающий и возвращающий именно этот тип. Когда конкретного типа нет и контейнер фактически равен Hold[Never], stored выставляется как обобщенный вызываемый объект def[R](R) -> R.

Почему это важно

Декораторы и утилиты высшего порядка, работающие с функциями T → T, в Python встречаются часто. Не понимая, как устроены области видимости переменных типов, легко ненароком стереть обобщенность и потерять типобезопасность или получить Never в ключевых местах. На практике проявляется разница между существованием и универсальной квантификацией: def (T) -> T для некоторого T — не то же самое, что def[T](T) -> T для любого T, и система типов проводит это различие на границах модуля.

Итоги

Ожидайте, что переменные типов связываются только внутри обобщенных классов и функций, а не на уровне модуля. Не рассчитывайте, что декоратор сохранит свободную переменную типа, когда вы превращаете обобщенную функцию в безаргументный thunk, возвращающий контейнер. Если нужно сохранить полезную типизацию при доступе к атрибуту, дескрипторный обходной путь позволяет развести два случая: когда конкретный тип известен и когда нет. Этот подход сознательно хаковый и может зависеть от текущих деталей вывода типов вроде Hold[Never]; он также не сохраняет границы и ограничения переменных типов. Если можете переработать API так, чтобы переменная типа оставалась связанной в списке параметров функции, вы полностью избегаете проблемы раннего связывания. В остальных случаях используйте дескрипторный трюк осознанно и документируйте компромиссы.