Using uuid.UUID with PotgreSQL's UUID

[bochecha]

I have a model using PostgreSQL's UUID:

from sqlalchemy.dialects import postgresql

class MyModel(Base):
    __tablename__ = 'my-model'

    some_id = Column(postgresql.UUID(as_uuid=True), nullable=False)

SQLAlchemy lets me create instances as follows:

import uuid

a_uuid = uuid.uuid4()

MyModel(some_id=a_uuid)
MyModel(some_id=str(a_uuid))

However, the stubs here seem to only accept the latter, not the former:

error: Incompatible type for "some_id" of "MyModel" (got "UUID", expected "str")

Could support for the former be added?

(I tested with sqlalchemy-stubs master)

[rafales]

@ilevkivskyi any idea how to approach this? SQLAlchemy uses this pattern all over the place, where the type depends on the passed arguments (eg. Numeric.as_decimal(), Enum, UUID). I guess a plugin will be needed to handle it.

[wbolster]

i also suffer from this issue.

is there a way to explicitly override the type that sqlalchemy-stubs infers for a column? that would at least avoid the false positives, at the (small) cost of putting an explicit type annotation in place.

i tried using annotations like sqlalchemy.Column[uuid.UUID] and sqlalchemy.types.TypeEngine[uuid.UUID] (also in string form) but no luck...

[wbolster]

UPDATE: there is a better way

it seems the following does the trick.

at the module level:

import uuid
from typing import cast, TYPE_CHECKING

import sqlalchemy
import sqlalchemy.dialects.postgresql

# https://github.com/dropbox/sqlalchemy-stubs/issues/94
_PostgreSQLUUID = sqlalchemy.dialects.postgresql.UUID(as_uuid=True)
if TYPE_CHECKING:
    PostgreSQLUUID = cast(sqlalchemy.types.TypeEngine[uuid.UUID], _PostgreSQLUUID)
else:
    PostgreSQLUUID = _PostgreSQLUUID

use in orm model classes then looks like this:

class Foo(Base):
    id = sqlalchemy.Column(PostgreSQLUUID, default=uuid.uuid4, primary_key=True)

[ilevkivskyi]

@wbolster Looks like you answered your own question. I think you can also just use # type: ignore on the definition with the explicit annotation. Also I think your way can be simplified to:

if TYPE_CHECKING:
    PostgreSQLUUID = sqlalchemy.types.TypeEngine[uuid.UUID]
else:
    PostgreSQLUUID = sqlalchemy.dialects.postgresql.UUID(as_uuid=True)

[wbolster]

hmm, not sure.

• reassigning to the same name with a different type will trigger another mypy error
• ignoring errors on the line defining the column does not suppress the errors elsewhere when constructing instances

[ilevkivskyi]

ignoring errors on the line defining the column does not suppress the errors elsewhere when constructing instances

This seems wrong, can you give an example?

[sirosen]

if TYPE_CHECKING:
    PostgreSQLUUID = sqlalchemy.types.TypeEngine[uuid.UUID]
else:
    PostgreSQLUUID = sqlalchemy.dialects.postgresql.UUID(as_uuid=True)

Could someone explain this workaround? The runtime value is obvious, but even after reading the source for TypeEngine, I can't tell what TypeEngine[uuid.UUID] is. It works for my codebase, but I'd like to be able to document what this is doing and why it satisfies mypy.

[wbolster]

try reveal_type(some_column) on a non-uuid column and observe the mypy output.

the type engine "trick" basically tells the sqlalchemy mypy plugin that this is a column that deals with uuid.UUID objects.

[sirosen]

reveal_type is handy, but I mean that I don't understand how sqlalchemy.types.TypeEngine[uuid.UUID] produces a Column[UUID*]. As far as I can tell, TypeEngine doesn't even support subscripting. So is this some magic hint to sqlmypy?

(By the way, thanks for sharing the workaround!)

[wbolster]

this is an update to my solution posted earlier. the good news is that things can be simplified. :partying_face:

the ‘if TYPE_CHECKING’ block and the workarounds using redefinitions can be avoided altogether by using a string literal for the type annotation (since the types are generic in stubs but not at runtime).

here's a full example; the :mage: trick here is the PostgreSQLUUID definition which can be used as the type for sqlalchemy.Column(...).

import uuid
from typing import cast

import sqlalchemy
import sqlalchemy.dialects.postgresql
import sqlalchemy.ext.declarative

PostgreSQLUUID = cast(
    "sqlalchemy.types.TypeEngine[uuid.UUID]",
    sqlalchemy.dialects.postgresql.UUID(as_uuid=True),
)

metadata = sqlalchemy.MetaData()
Base = sqlalchemy.ext.declarative.declarative_base(metadata=metadata)

class SomeClass(Base):
    id = sqlalchemy.Column(PostgreSQLUUID, default=uuid.uuid4, primary_key=True)
    ... # other columns go here

[wbolster]

it can even be made more strict by using a NewType for a class-specific id type that cannot be silently conflated with other uuids (without being explicit about it):

from typing import NewType

FooId = NewType("FooId", uuid.UUID)
PostgreSQLFooId = cast("sqlalchemy.types.TypeEngine[FooId]", PostgreSQLUUID)

class Foo(Base):
    id = sqlalchemy.Column(PostgreSQLFooId, default=uuid.uuid4, primary_key=True)

usage example:

obj = session.query(Foo).one()
reveal_type(obj.id)  # FooID, not uuid.UUID

this means it's possible to have different AddressId and UserId types, even if both are the same underlying type (such as an integer or uuid)!

see https://twitter.com/wbolster/status/1299319516829298689

[rafales]

The most ergonomic way to do this (although not in libraries, strictly in your own project code) is just to monkey-patch classes to add runtime support for Class[] syntax. This is what django-stubs project is doing internally for QuerySet and Manager: https://github.com/typeddjango/django-stubs/blob/master/mypy_django_plugin/django/context.py#L60