How can I type hint SQLAlchemy Queries?

Question

update 2024: sqlalchemy 2.0 has much better typing support

I'm using Pylance (with type hints powered by PyRight) in VSCode, and have the following function:

def sample(session: sqlalchemy.orm.Session) -> sqlalchemy.orm.query.Query:
    return session.query((MyModel)

When looking at Pylance's type inferences, this function actually has the return type of Query[Unknown]. When I change the type hint to sqlalchemy.orm.query.Query[MyModel], Pylance correctly can deduce relevant types for it. this is great! Until I actually run the code, and get TypeError: 'type' object is not subscriptable.

It makes sense why I get this error, but it seems like there must be some way to make this work, otherwise Pylance is just taunting me. One workaround (which does work) I have found is this:

def sample(session: sqlalchemy.orm.Session) -> (
    sqlalchemy.orm.query.Query[MyModel]
    if typing.TYPE_CHECKING
    else sqlalchemy.orm.query.Query
):
    return session.query(MyModel)

but this is obviously awful and I hate it. I might be able to make a helper function that extracts that and is more readable, but I really feel like this should be possible without an awful workaround!

Flaw

Doc for type hint support in SQLA 1.4 docs.sqlalchemy.org/en/14/orm/extensions/mypy.html – Mayhap

Answer 1

def sample(session: sqlalchemy.orm.Session) -> 
"sqlalchemy.orm.query.Query[MyModel]":
    return session.query(MyModel)

actually works just fine and I feel dumb. If there is any other way I'd be happy to accept that answer!

Flaw