Plugin: handle declarative_base cls argument

sqlmypy.py

@@ -1,15 +1,16 @@
+from mypy.mro import calculate_mro, MroError
 from mypy.plugin import Plugin, FunctionContext, ClassDefContext
 from mypy.plugins.common import add_method
-from mypy.nodes import(
+from mypy.nodes import (
     NameExpr, Expression, StrExpr, TypeInfo, ClassDef, Block, SymbolTable, SymbolTableNode, GDEF,
-    Argument, Var, ARG_STAR2, MDEF
+    Argument, Var, ARG_STAR2, MDEF, TupleExpr, RefExpr
 )
 from mypy.types import (
     UnionType, NoneTyp, Instance, Type, AnyType, TypeOfAny, UninhabitedType, CallableType
 )
 from mypy.typevars import fill_typevars_with_any
 
-from typing import Optional, Callable, Dict, TYPE_CHECKING
+from typing import Optional, Callable, Dict, TYPE_CHECKING, List
 if TYPE_CHECKING:
     from typing_extensions import Final
 
@@ -141,14 +142,37 @@ def decl_info_hook(ctx):
 
         Base = declarative_base()
     """
+    cls_bases = []  # type: List[Instance]
+
+    # Passing base classes as positional arguments is currently not handled.
+    if 'cls' in ctx.call.arg_names:
+        declarative_base_cls_arg = ctx.call.args[ctx.call.arg_names.index("cls")]
+        if isinstance(declarative_base_cls_arg, TupleExpr):
+            items = [item for item in declarative_base_cls_arg.items]
+        else:
+            items = [declarative_base_cls_arg]
+
+        for item in items:
+            if isinstance(item, RefExpr) and isinstance(item.node, TypeInfo):
+                base = fill_typevars_with_any(item.node)
+                # TODO: Support tuple types?
+                if isinstance(base, Instance):
+                    cls_bases.append(base)
+
     class_def = ClassDef(ctx.name, Block([]))
     class_def.fullname = ctx.api.qualified_name(ctx.name)
 
     info = TypeInfo(SymbolTable(), class_def, ctx.api.cur_mod_id)
     class_def.info = info
     obj = ctx.api.builtin_type('builtins.object')
-    info.mro = [info, obj.type]
-    info.bases = [obj]
+    info.bases = cls_bases or [obj]
+    try:
+        calculate_mro(info)
+    except MroError:
+        ctx.api.fail("Not able to calculate MRO for declarative base", ctx.call)
+        info.bases = [obj]
+        info.fallback_to_any = True
+
     ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info))
     set_declarative(info)
 

test/test-data/sqlalchemy-plugin-features.test

@@ -219,3 +219,57 @@ class User(Base):
 record = {'name': 'John Doe'}
 User(**record)  # OK
 [out]
+
+[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]