sqlalchemy-plugin-features.test

[case testModelInitColumnDeclared]
# flags: --strict-optional
from sqlalchemy import Column, Integer, String
from base import Base
from typing import Optional

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer(), primary_key=True)
    name = Column(String())

oi: Optional[int]
os: Optional[str]

User()
User(1, 2)  # E: Too many arguments for "User"
User(id=int(), name=str())
User(id=oi)  # E: Incompatible type for "id" of "User" (got "Optional[int]", expected "int")
User(name=os)

[file base.py]
from sqlalchemy.ext.declarative import declarative_base
Base = declarative_base()
[out]

[case testModelInitColumnDecorated]
# flags: --strict-optional
from sqlalchemy import Column, Integer, String
from base import Base
from typing import Optional

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer(), primary_key=True)
    name = Column(String())

oi: Optional[int]
os: Optional[str]

User()
User(1, 2)  # E: Too many arguments for "User"
User(id=int(), name=str())
User(id=oi)  # E: Incompatible type for "id" of "User" (got "Optional[int]", expected "int")
User(name=os)

[file base.py]
from sqlalchemy.ext.declarative import as_declarative
@as_declarative()
class Base:
    ...
[out]

[case testModelInitRelationship]
from typing import TYPE_CHECKING, List

from sqlalchemy import Column, Integer, String
from sqlalchemy.orm import relationship, RelationshipProperty
from base import Base
if TYPE_CHECKING:
    from other import Other

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer(), primary_key=True)
    other = relationship('Other', uselist=False)
    many_others = relationship(Other, uselist=True)

o: Other
mo: List[Other]
User()
User(other=o, many_others=mo)
User(other=mo)  # E: Incompatible type for "other" of "User" (got "List[Other]", expected "Other")
User(unknown=42)  # E: Unexpected column "unknown" for model "User"

[file other.py]
from sqlalchemy import Column, Integer, String
from base import Base

class Other(Base):
    __tablename__ = 'other'
    id = Column(Integer(), primary_key=True)
    name = Column(String(), nullable=False)

[file base.py]
from sqlalchemy.ext.declarative import declarative_base
Base = declarative_base()
[out]

[case testRelationshipType]
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import Column, Integer, String
from sqlalchemy.orm import relationship, RelationshipProperty
from typing import Iterable

Base = declarative_base()

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer(), primary_key=True)
    first_other = relationship(Other)
    second_other = relationship(Other, uselist=True)
    bad_other: RelationshipProperty[int] = relationship(Other, uselist=False)  # E: Incompatible types in assignment (expression has type "RelationshipProperty[Other]", variable has type "RelationshipProperty[int]")

user = User()
reveal_type(user.first_other)  # E: Revealed type is 'main.Other*'
reveal_type(user.second_other)  # E: Revealed type is 'typing.Iterable*[main.Other]'

class Other(Base):
    __tablename__ = 'other'
    id = Column(Integer(), primary_key=True)
[out]

[case testRelationshipString]
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import Column, Integer, String
from sqlalchemy.orm import relationship, RelationshipProperty
from typing import Iterable

Base = declarative_base()

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer(), primary_key=True)
    first_other = relationship('Other', uselist=False)
    second_other = relationship('Other', uselist=True)
    second_bad_other = relationship('What')  # E: Cannot find model "What" \
                                             # N: Only imported models can be found; use "if TYPE_CHECKING: ..." to avoid import cycles

user = User()
reveal_type(user.first_other)  # E: Revealed type is 'main.Other*'
reveal_type(user.second_other)  # E: Revealed type is 'typing.Iterable*[main.Other]'

class Other(Base):
    __tablename__ = 'other'
    id = Column(Integer(), primary_key=True)
[out]

[case testRelationshipAnnotated]
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import Column, Integer, String
from sqlalchemy.orm import relationship, RelationshipProperty
from typing import Iterable

Base = declarative_base()

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer(), primary_key=True)
    first_other: RelationshipProperty[Other] = relationship('Other')
    second_other: RelationshipProperty[Iterable[Other]] = relationship(Other, uselist=True)
    third_other: RelationshipProperty[Other] = relationship(Other, uselist=False)
    bad_other: RelationshipProperty[Other] = relationship('Other', uselist=True)  # E: Incompatible types in assignment (expression has type "RelationshipProperty[Iterable[Other]]", variable has type "RelationshipProperty[Other]")

class Other(Base):
    __tablename__ = 'other'
    id = Column(Integer(), primary_key=True)
[out]

[case testColumnCombo]
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import Column, Integer, String

Base = declarative_base()

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer(), primary_key=True)
    name = Column(String(), default='John Doe', nullable=True)

user: User
reveal_type(User.name)  # E: Revealed type is 'sqlalchemy.sql.schema.Column[Union[builtins.str*, None]]'
[out]

[case testAddedAttributesDeclared]
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import Column, Integer, String

Base = declarative_base()

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer(), primary_key=True)
    name = Column(String(), default='John Doe', nullable=True)

user: User
reveal_type(User.metadata)  # E: Revealed type is 'sqlalchemy.sql.schema.MetaData'
reveal_type(User.__table__)  # E: Revealed type is 'sqlalchemy.sql.schema.Table'
[out]

[case testAddedAttributedDecorated]
from sqlalchemy.ext.declarative import as_declarative
from sqlalchemy import Column, Integer, String

@as_declarative()
class Base:
    ...

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer(), primary_key=True)
    name = Column(String(), default='John Doe', nullable=True)

user: User
reveal_type(User.metadata)  # E: Revealed type is 'sqlalchemy.sql.schema.MetaData'
reveal_type(User.__table__)  # E: Revealed type is 'sqlalchemy.sql.schema.Table'
[out]

[case testKwArgsModelOK]
from sqlalchemy import Column, Integer, String
from sqlalchemy.ext.declarative import declarative_base

Base = declarative_base()

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True)
    name = Column(String)

record = {'name': 'John Doe'}
User(**record)  # OK
[out]

[case testGrouping]
from sqlalchemy import Column, String, text
from sqlalchemy.sql.elements import Grouping

reveal_type(Grouping(Column(String, nullable=False)))  # E: Revealed type is 'sqlalchemy.sql.elements.Grouping[builtins.str*]'
reveal_type(Grouping(text('foo')))  # E: Revealed type is 'sqlalchemy.sql.elements.Grouping[None]'

[case testDeclarativeBaseWithBaseClass]
from sqlalchemy import Column, Integer, String
from base import Base

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer(), primary_key=True)
    name = Column(String())

user: User
reveal_type(user.f())  # E: Revealed type is 'builtins.str'

[file base.py]
from sqlalchemy.ext.declarative import declarative_base

class Model:
    def f(self) -> str: ...
Base = declarative_base(cls=Model)
[out]

[case testDeclarativeBaseWithMultipleBaseClasses]
from sqlalchemy import Column, Integer, String
from base import Base

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer(), primary_key=True)
    name = Column(String())

user: User
reveal_type(user.f())  # E: Revealed type is 'builtins.str'
reveal_type(user.g())  # E: Revealed type is 'builtins.int'

[file base.py]
from sqlalchemy.ext.declarative import declarative_base

class Model:
    def f(self) -> str: ...
class Model2:
    def g(self) -> int: ...
Base = declarative_base(cls=(Model, Model2))
[out]

[case testDeclarativeBaseWithBaseClassWrongMRO]
from sqlalchemy.ext.declarative import declarative_base

class M1:
    ...
class M2(M1):
    ...
Base = declarative_base(cls=(M1, M2))  # E: Not able to calculate MRO for declarative base
reveal_type(Base)  # E: Revealed type is 'Any'
[out]

[case testRelationshipIsGuessed]
from sqlalchemy import Column, Integer, String, ForeignKey
from sqlalchemy.orm import relationship
from sqlalchemy.ext.declarative import declarative_base

Base = declarative_base()

class Parent(Base):
    __tablename__ = 'parents'
    id = Column(Integer, primary_key=True)
    name = Column(String)

    children = relationship("Child")

class Child(Base):
    __tablename__ = 'children'
    id = Column(Integer, primary_key=True)
    name = Column(String)
    parent_id = Column(Integer, ForeignKey(Parent.id))

    parent = relationship(Parent)

child: Child
parent: Parent

reveal_type(child.parent)  # E: Revealed type is 'main.Parent*'
reveal_type(parent.children)  # E: Revealed type is 'typing.Iterable*[main.Child]'

[out]

[case testRelationshipIsGuessed2]
from sqlalchemy import Column, Integer, String, ForeignKey
from sqlalchemy.orm import relationship
from sqlalchemy.ext.declarative import declarative_base

Base = declarative_base()

class Parent(Base):
    __tablename__ = 'parents'
    id = Column(Integer, primary_key=True)
    name = Column(String)

    children = relationship("Child")

class Child(Base):
    __tablename__ = 'children'
    id = Column(Integer, primary_key=True)
    name = Column(String)
    parent_id = Column(Integer, ForeignKey("parents.id"))

    parent = relationship(Parent)

child: Child
parent: Parent

reveal_type(child.parent)  # E: Revealed type is 'main.Parent*'
reveal_type(parent.children)  # E: Revealed type is 'typing.Iterable*[main.Child]'

[out]

[case testRelationshipIsGuessed3]
from sqlalchemy import Column, Integer, String, ForeignKey
from sqlalchemy.orm import relationship
from sqlalchemy.ext.declarative import declarative_base

Base = declarative_base()

class Parent(Base):
    __tablename__ = 'parents'
    id = Column(Integer, primary_key=True)
    name = Column(String)

    children = relationship("Child")

class Child(Base):
    __tablename__ = 'children'
    id = Column(Integer, primary_key=True)
    name = Column(String)
    parent_id = Column(Integer, ForeignKey("other_parents.id"))

    parent = relationship(Parent)

child: Child
parent: Parent

reveal_type(child.parent)  # E: Revealed type is 'main.Parent*'
reveal_type(parent.children)  # E: Revealed type is 'main.Child*'

[out]