fastapi / sqlmodel

SQL databases in Python, designed for simplicity, compatibility, and robustness.
https://sqlmodel.tiangolo.com/
MIT License
14.44k stars 657 forks source link

Pyright type checking when working with aggregate data #450

Open Bryley opened 2 years ago

Bryley commented 2 years ago

First Check

Commit to Help

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()

pyright1 pyright2

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 commented 2 years ago

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 commented 2 years ago

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']}")