Mypy error when inferring type of Enum column

[yonaszmerla]

First Check

• I added a very descriptive title here.
• I used the GitHub search to find a similar question 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

import enum
from typing import Optional
from sqlmodel import Field, SQLModel, Session, create_engine, select
import sqlalchemy as sa


class TestEnum(str, enum.Enum):
    TEST_1 = "test_1"
    TEST_2 = "test_2"
    
    
class MyTest(SQLModel, table=True):
    __tablename__ = "my_test"
    
    id: Optional[int] = Field(None ,primary_key=True)
    test_enum: Optional[TestEnum] = Field(sa_column=sa.Column(sa.Enum(TestEnum), nullable=True))


engine = create_engine("sqlite:///test.db", echo=True)

SQLModel.metadata.create_all(engine)

session = Session(engine)

session.add(MyTest(test_enum=TestEnum.TEST_1))
session.commit()

def get_test_enum() -> Optional[TestEnum]:
    return session.exec(
        select(MyTest.test_enum)
    ).first()

test_obj_enum = get_test_enum()

print(type(test_obj_enum))
print(test_obj_enum)

Description

Create MyTest SQLModel class with custom str_enum field
Retrieve this column from the DB using sqlmodel
The type inferred by mypy is str and not enum as it should be

I will get an error from MyPy

error: Incompatible return value type (got "str | None", expected "TestEnum | None")  [return-value]
Found 1 error in 1 file (checked 1 source file)

On the return value of def get_test_enum() -> Optional[TestEnum]:

Operating System

macOS

Operating System Details

Ventura 13.2.1

SQLModel Version

0.0.11

Python Version

3.9.18

Additional Context

No response

[YuriiMotov]

Using col() resolves the error:

from sqlmodel import col

...

def get_test_enum() -> Optional[TestEnum]:
    return session.exec(
        select(col(MyTest.test_enum))
    ).first()