sqlalchemy-stubs icon indicating copy to clipboard operation
sqlalchemy-stubs copied to clipboard

Enum interpreted as str

Open qsantos opened this issue 5 years ago • 10 comments

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

qsantos avatar Oct 02 '19 08:10 qsantos

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

ilevkivskyi avatar Oct 14 '19 16:10 ilevkivskyi

I see, thanks for the information.

qsantos avatar Oct 14 '19 17:10 qsantos

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

DeanWay avatar Nov 04 '19 15:11 DeanWay

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 avatar Nov 04 '19 15:11 ilevkivskyi

@ilevkivskyi Thank you! That seems to fix the problem for me

DeanWay avatar Nov 05 '19 15:11 DeanWay

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 avatar Jan 07 '20 16:01 helgridly

@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

DeanWay avatar Jan 07 '20 16:01 DeanWay

Thank you!

helgridly avatar Jan 07 '20 19:01 helgridly

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.

cardoe avatar Jan 22 '20 17:01 cardoe

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]

vincentwyshan avatar May 04 '20 07:05 vincentwyshan