sqlalchemy-stubs
sqlalchemy-stubs copied to clipboard
Enum interpreted as str
I have an SQLAlchemy model that makes use of the Enum
column type. When accessing the field of an instance of this model, mypy believes that the type of the field is str
even though it is actually an enum (e.g. MyEnum
). This is annoying since, when I do want to access its value, mypy fails with error: "str" has no attribute "value"
.
Although it cannot be run as such, the following snippet demonstrates the behavior when run through mypy. I would expect mypy to expect m.state
should be of type MyEnum
(really, in this snippet, it will be None
, but in real code, it will be a MyEnum
).
import enum
from sqlalchemy import Column, Enum
from sqlalchemy.ext.declarative import declarative_base
class MyEnum(enum.Enum):
A = 'A'
B = 'B'
Base = declarative_base()
class MyModel(Base):
state = Column(Enum(MyEnum), nullable=False)
m = MyModel()
m.state.value
This is because SQLAlchemy enum predates PEP 435, so it is actually possible to pass a list of strings to the Enum
constructor.
We can try to support this by making Enum
stub definition https://github.com/dropbox/sqlalchemy-stubs/blob/master/sqlalchemy-stubs/sql/sqltypes.pyi#L170 generic and using an overloaded constructor to bind the type argument (note however overloaded generic constructors are only supported in latest dev version of mypy).
I see, thanks for the information.
Is there a work around here to allow the type checker to treat this as an enum? Currently I'm just using # type: ignore
on lines that use an Enum column
Something like this should work I think:
if TYPE_CHECKING:
from sqlalchemy.sql.type_api import TypeEngine
class Enum(TypeEngine[T]):
def __init__(self, enum: Type[T]) -> None: ...
else:
from sqlalchemy import Enum
@ilevkivskyi Thank you! That seems to fix the problem for me
I've reached this thread from Google and am trying to implement the answer @ilevkivskyi provided.
I've figured out that it requires from typing import TYPE_CHECKING, TypeEngine
, but I don't know what to do with T
. I can get it working for one enum class by replacing T
with the enum class name, but what if I have more than one enum type?
@helgridly T here is a TypeVar, Type
and TypeVar
can be imported from typing.
Altogether like this:
from typing import TypeVar, Type, TYPE_CHECKING
if TYPE_CHECKING:
from sqlalchemy.sql.type_api import TypeEngine
T = TypeVar('T')
class Enum(TypeEngine[T]):
def __init__(self, enum: Type[T]) -> None: ...
else:
from sqlalchemy import Enum
Thank you!
Still not 100% correct. Enum can take some kwargs.
from typing import Any, TypeVar, Type, TYPE_CHECKING
if TYPE_CHECKING:
from sqlalchemy.sql.type_api import TypeEngine
T = TypeVar('T')
class Enum(TypeEngine[T]):
def __init__(self, enum: Type[T], **kwargs: Any) -> None: ...
else:
from sqlalchemy import Enum
Probably would be better to make all the possible kwargs checked.
Still not 100% correct. Enum can take some kwargs.
from typing import Any, TypeVar, Type, TYPE_CHECKING if TYPE_CHECKING: from sqlalchemy.sql.type_api import TypeEngine T = TypeVar('T') class Enum(TypeEngine[T]): def __init__(self, enum: Type[T], **kwargs: Any) -> None: ... else: from sqlalchemy import Enum
Probably would be better to make all the possible kwargs checked.
Tried this solution but still get error
state = Column(Enum(MyEnum), nullable=False)
# error: Need type annotation for 'state' [var-annotated]