Как типизировать ORM‑классы SQLAlchemy с целочисленным id в mypy: Protocol и ClassVar

Как типизировать ORM‑класс SQLAlchemy с целочисленным id в mypy: Protocol с id: ClassVar[orm.Mapped[int]] для select/where/order_by без Any и строгой типизации.

Статическая типизация и SQLAlchemy часто пересекаются на непростом участке: атрибуты, сопоставленные на уровне класса, ведут себя по‑разному на самом классе и на экземпляре. Частая задача — типизировать функцию, которая принимает любой ORM‑класс SQLAlchemy с целочисленным первичным ключом id. Цель — сохранить строгую типизацию и не скатываться к id: Any, при этом поддерживая выражения вроде model.id > 42 и model.id.desc() в цепочке select/where/order_by. Окружение: Python 3.12.3, SQLAlchemy 2.0.41 и проверка типов mypy.

Problem

Задача — написать функцию, которая принимает любой ORM‑класс с целочисленным id и возвращает типизированный SQLAlchemy Select. Несколько попыток описать ограничение через Protocol и TypeVar приводили к ошибкам mypy, из‑за чего приходилось идти на компромисс с id: Any.

from typing import Any, Protocol, TypeVar
import sqlalchemy as sa
from sqlalchemy import orm


class WithIntPk(Protocol):
    id: Any  # weaker than desired, but type checks

U = TypeVar('U', bound=WithIntPk)

def build_query(model_cls: type[U]) -> sa.Select[tuple[U]]:
    return sa.select(model_cls).where(model_cls.id > 42).order_by(model_cls.id.desc())


class OrmBase(orm.DeclarativeBase):
    pass

class Widget(OrmBase):
    __tablename__ = 'widgets'
    id: orm.Mapped[int] = orm.mapped_column(primary_key=True)

print(build_query(Widget))

Ранее пробовали id: int, id: orm.Mapped[int] и id: orm.attributes.InstrumentedAttribute[int]. mypy отвечал ошибками вроде “Argument 1 to 'where' of 'Select' has incompatible type 'bool' ... [arg-type]” и “Value of type variable 'T' of 'f' cannot be 'Foo' [type-var]”.

Why this happens

Суть в том, что поля SQLAlchemy на уровне класса возвращают SQL‑выражения и поддерживают операторы и методы, которых нет у обычных целых чисел. Иными словами, model.id > 42 — это не булево значение, а SQLAlchemy‑выражение. Точно так же model.id.desc() — это метод объекта SQL‑выражения, а не int. Поэтому аннотация как простой int не подходит для использования в where и order_by.

Есть и тонкость с Protocol: атрибуты, объявленные в Protocol, по умолчанию описывают атрибуты экземпляра, тогда как функция работает с атрибутами класса (model_cls.id). Сама SQLAlchemy предоставляет разные типы в зависимости от того, обращаетесь ли вы к id на классе или на экземпляре. Разницу видно при reveal_type:

reveal_type(Foo.id) — “sqlalchemy.sql.elements.ColumnElement[builtins.bool]”
reveal_type(Foo().id) — “builtins.int”

Значит, Protocol должен описывать атрибут класса с типом, соответствующим сопоставленному атрибуту SQLAlchemy, а не атрибут экземпляра.

Solution

Решение — указать для атрибута тип orm.Mapped[int] и обернуть его в ClassVar, чтобы зафиксировать контракт на уровне класса. Это отражает то, как типизация SQLAlchemy представляет атрибут: на классе он ведёт себя как сопоставленный, поддерживающий выражения атрибут; на экземплярах сводится к базовому типу Python.

from typing import ClassVar, Protocol, TypeVar
import sqlalchemy as sa
from sqlalchemy import orm


class IntPkContract(Protocol):
    id: ClassVar[orm.Mapped[int]]

R = TypeVar('R', bound=IntPkContract)

def make_select(entity: type[R]) -> sa.Select[tuple[R]]:
    return sa.select(entity).where(entity.id > 42).order_by(entity.id.desc())


class BaseModel(orm.DeclarativeBase):
    pass

class Thing(BaseModel):
    __tablename__ = 'things'
    id: orm.Mapped[int] = orm.mapped_column(primary_key=True)

print(make_select(Thing))

Так сохраняется точная типизация и при этом остаётся привычное использование выражений SQLAlchemy в where и order_by.

What to keep in mind

orm.Mapped[int] — подходящий целевой тип для атрибута id в Protocol, потому что он поддерживает перегрузки операторов и методы, используемые в запросах SQLAlchemy. ClassVar обеспечивает, что вы задаёте контракт именно для атрибута класса, а не экземпляра. На экземплярах SQLAlchemy по‑прежнему сводит атрибут к int, что соответствует ожидаемому поведению.

Также отмечают, что простой id: orm.Mapped[int] в Protocol проходит проверку в pyright, тогда как для mypy нужен ClassVar. Есть мнение, что это может быть багом mypy; приведены ссылка на issue mypy и обсуждение в pyright. Если поведение в будущем изменится, необходимость в ClassVar в этом сценарии может пересмотреться. References: https://github.com/python/mypy/issues/19702 и https://github.com/microsoft/pyright/discussions/7942.

Why this matters

Точная типизация моделей SQLAlchemy приносит ощутимую пользу: защищает от тихих регрессий, упрощает безопасный рефакторинг и заставляет код построения запросов быть честным относительно формы и поведения сопоставленных атрибутов. Отказ от id: Any не даёт целым цепочкам выражений провалиться в нетипизированную область, что может скрыть тонкие ошибки.

Takeaways

Если вам нужна функция, принимающая любой ORM‑класс SQLAlchemy с целочисленным id и вы используете mypy, объявите Protocol с id: ClassVar[orm.Mapped[int]] и привяжите к нему TypeVar. Это удовлетворяет и доступу на уровне класса, и семантике выражений, требуемой select, where и order_by. Следите за поведением проверяющих типов в mypy и pyright и по возможности предпочитайте точные аннотации вместо Any. Справка по типизации атрибутов класса: документация Python по typing.ClassVar — https://docs.python.org/3/library/typing.html#typing.ClassVar.