Pyright type checking when working with aggregate data

[Bryley]

First Check

• I added a very descriptive title to this issue.
• I used the GitHub search to find a similar issue and didn't find it.
• I searched the SQLModel documentation, with the integrated search.
• I already searched in Google "How to X in SQLModel" and didn't find any information.
• I already read and followed all the tutorial in the docs and didn't find an answer.
• I already checked if it is not related to SQLModel but to Pydantic.
• I already checked if it is not related to SQLModel but to SQLAlchemy.

Commit to Help

• I commit to help with one of those options 👆

Example Code

from typing import Optional
from sqlmodel import Field, SQLModel, Session, create_engine, func, select

class Score(SQLModel, table=True):
    __tablename__ = "scores" # type: ignore

    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    score: int = Field(ge=0, le=100)

engine = create_engine("sqlite:///")
SQLModel.metadata.create_all(engine)

with Session(engine) as session:
    session.add(Score(name="Bob", score=51))
    session.add(Score(name="Alice", score=88))
    session.add(Score(name="Greg", score=12))
    session.add(Score(name="Patrick", score=100))
    session.add(Score(name="Susie", score=45))

    session.commit()

with Session(engine) as session:
    res = session.exec(
        select(
            func.min(Score.score),
            func.max(Score.score)
        )
    ).one()

    print(f"Minimum score: {res['min']}")
    print(f"Maximum score: {res['max']}")

Description

The above code works as expected and outputs:

Minimum score: 12
Maximum score: 100

However the issue I'm having is the pyright static type checking around:

res = session.exec(
    select(
        func.min(Score.score),
        func.max(Score.score)
    )
).one()

Pyright gives me the following error and the res variable has an Unkown type according to Pyright.

Argument of type "Function[NullType]" cannot be assigned to parameter "entity_1" of type "Type[_TModel_1@select]" in function "select"
    Type "Function[NullType]" cannot be assigned to type "Type[_TModel_1@select]

Expected Result / Solution Idea

The res variable should not have an unknown type.
Given the following code, the type of res should be something like CustomMinMax where CustomMinMax is a generated pydantic model/named tuple/dataclass like the following:

class CustomMinMax(BaseModel):
    min: int
    max: int

Not sure how possible that would be to do.

Has anybody faced the same issue or has any ideas on ways around this so the res variable has an appropriate type, or maybe I'm just missing a pyright config option that solves this problem?

Operating System

Linux

Operating System Details

Arch

SQLModel Version

0.0.8

Python Version

3.9.13

Additional Context

Pyright version: 1.1.271

[meirdev]

It's not a perfect solution, but you can use declared_attr and column_property:

from typing import cast
from sqlalchemy.orm import column_property, declared_attr

...

class ScoreTableBase(SQLModel):
    __tablename__ = "scores" # type: ignore

    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    score: int = Field(ge=0, le=100)


class Score(ScoreTableBase, table=True):
    pass


class ScoreAggregate(ScoreTableBase, table=True):
    __table_args__ = {"extend_existing": True}

    @declared_attr
    def max(cls) -> int:
        return cast(int, column_property(func.max(cls.score)))

    @declared_attr
    def min(cls) -> int:
        return cast(int, column_property(func.min(cls.score)))

...

with Session(engine) as session:
    res = session.exec(
        select(
            ScoreAggregate
        )
    ).one()

    print(f"Minimum score: {res.min}")
    print(f"Maximum score: {res.max}")

[phi-friday]

use
Session.exec -> Session.execute
sqlmodel.select -> sqlalchemy.select

ex:

from typing import ClassVar, Optional

from sqlalchemy import select
from sqlmodel import Field, Session, SQLModel, create_engine, func


class Score(SQLModel, table=True):
    __tablename__: ClassVar[str] = "scores"

    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    score: int = Field(ge=0, le=100)


engine = create_engine("sqlite:///")
SQLModel.metadata.create_all(engine)

with Session(engine) as session:
    session.add(Score(name="Bob", score=51))
    session.add(Score(name="Alice", score=88))
    session.add(Score(name="Greg", score=12))
    session.add(Score(name="Patrick", score=100))
    session.add(Score(name="Susie", score=45))

    session.commit()

with Session(engine) as session:
    res = session.execute(select(func.min(Score.score), func.max(Score.score))).one()

    print(f"Minimum score: {res['min']}")
    print(f"Maximum score: {res['max']}")