Infer type for dialects.postgresql.UUID(as_uuid=True)

Is your feature request related to a problem? Please describe.

from sqlalchemy import Column, ForeignKey
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.orm import declarative_base

Base = declarative_base()
UUIDC = UUID(as_uuid=True)


class Foo(Base):
    id = Column(UUIDC)
    bar_id = Column(UUIDC, ForeignKey("bars.id"))

This will yield the following errors:

$ mypy --strict testuuidc.py
testuuidc.py:10: error: Need type annotation for 'id'
testuuidc.py:11: error: [SQLAlchemy Mypy plugin] Can't infer type from ORM mapped expression assigned to attribute 'bar_id'; please specify a Python type or Mapped[<python type>] on the left hand side.
Found 2 errors in 1 file (checked 1 source file)

Describe the solution you’d like

I wonder if thanks to the plugin, mypy wouldn’t be able to infer the uuid.UUID type from the postgresql.UUID(as_uuid=True) definition.

Describe alternatives you’ve considered

Obviously, we can always provide the typing explicitly.

Have a nice day!

Thanks a lot for the amazing work on sqlalchemy!

[charlax]

Actually, providing the typing explictly fails:

from uuid import UUID

from sqlalchemy import Column, ForeignKey, Integer
from sqlalchemy.dialects.postgresql import UUID as _UUIDC
from sqlalchemy.orm import declarative_base, Mapped

Base = declarative_base()
UUIDC = _UUIDC(as_uuid=True)


class Foo(Base):
    id = Column(UUIDC)
    id2: UUID = Column(UUIDC)
    id3: Mapped[UUID] = Column(UUIDC)

    bar_id = Column(UUIDC, ForeignKey("bars.id"))

    int1 = Column(Integer)
    int2: int = Column(Integer)
    int3: Mapped[int] = Column(Integer)

This yields the following errors:

$ mypy --strict testuuidc.py
testuuidc.py:12: error: Need type annotation for 'id'
testuuidc.py:13: error: Incompatible types in assignment (expression has type "Column[<nothing>]", variable has type "UUID")
testuuidc.py:14: error: Incompatible types in assignment (expression has type "Column[<nothing>]", variable has type "Mapped[UUID]")
testuuidc.py:16: error: [SQLAlchemy Mypy plugin] Can't infer type from ORM mapped expression assigned to attribute 'bar_id'; please specify a Python type or Mapped[<python type>] on the left hand side.
Found 4 errors in 1 file (checked 1 source file)

[sqla-tester]

Mike Bayer has proposed a fix for this issue in the main branch:

add typing for PG UUID, other types https://gerrit.sqlalchemy.org/c/sqlalchemy/sqlalchemy/+/3870