Update mypy plugin to conform to strict mode

Change-Id: I09a3df5af2f2d4ee34d8d72c3dedc4f236df8eb1

1 files changed, 77 insertions, 45 deletions

diff --git a/lib/sqlalchemy/ext/mypy/names.py b/lib/sqlalchemy/ext/mypy/names.py
index 174a8f422..6ee600cd7 100644
--- a/lib/sqlalchemy/ext/mypy/names.py
+++ b/lib/sqlalchemy/ext/mypy/names.py
@@ -5,40 +5,48 @@
 # This module is part of SQLAlchemy and is released under
 # the MIT License: http://www.opensource.org/licenses/mit-license.php
 
+from typing import Dict
 from typing import List
+from typing import Optional
+from typing import Set
+from typing import Tuple
+from typing import Union
 
 from mypy.nodes import ClassDef
 from mypy.nodes import Expression
 from mypy.nodes import FuncDef
-from mypy.nodes import RefExpr
+from mypy.nodes import MemberExpr
+from mypy.nodes import NameExpr
 from mypy.nodes import SymbolNode
 from mypy.nodes import TypeAlias
 from mypy.nodes import TypeInfo
-from mypy.nodes import Union
 from mypy.plugin import SemanticAnalyzerPluginInterface
+from mypy.types import CallableType
+from mypy.types import get_proper_type
+from mypy.types import Instance
 from mypy.types import UnboundType
 
 from ... import util
 
-COLUMN = util.symbol("COLUMN")
-RELATIONSHIP = util.symbol("RELATIONSHIP")
-REGISTRY = util.symbol("REGISTRY")
-COLUMN_PROPERTY = util.symbol("COLUMN_PROPERTY")
-TYPEENGINE = util.symbol("TYPEENGNE")
-MAPPED = util.symbol("MAPPED")
-DECLARATIVE_BASE = util.symbol("DECLARATIVE_BASE")
-DECLARATIVE_META = util.symbol("DECLARATIVE_META")
-MAPPED_DECORATOR = util.symbol("MAPPED_DECORATOR")
-COLUMN_PROPERTY = util.symbol("COLUMN_PROPERTY")
-SYNONYM_PROPERTY = util.symbol("SYNONYM_PROPERTY")
-COMPOSITE_PROPERTY = util.symbol("COMPOSITE_PROPERTY")
-DECLARED_ATTR = util.symbol("DECLARED_ATTR")
-MAPPER_PROPERTY = util.symbol("MAPPER_PROPERTY")
-AS_DECLARATIVE = util.symbol("AS_DECLARATIVE")
-AS_DECLARATIVE_BASE = util.symbol("AS_DECLARATIVE_BASE")
-DECLARATIVE_MIXIN = util.symbol("DECLARATIVE_MIXIN")
-
-_lookup = {
+COLUMN: int = util.symbol("COLUMN")  # type: ignore
+RELATIONSHIP: int = util.symbol("RELATIONSHIP")  # type: ignore
+REGISTRY: int = util.symbol("REGISTRY")  # type: ignore
+COLUMN_PROPERTY: int = util.symbol("COLUMN_PROPERTY")  # type: ignore
+TYPEENGINE: int = util.symbol("TYPEENGNE")  # type: ignore
+MAPPED: int = util.symbol("MAPPED")  # type: ignore
+DECLARATIVE_BASE: int = util.symbol("DECLARATIVE_BASE")  # type: ignore
+DECLARATIVE_META: int = util.symbol("DECLARATIVE_META")  # type: ignore
+MAPPED_DECORATOR: int = util.symbol("MAPPED_DECORATOR")  # type: ignore
+COLUMN_PROPERTY: int = util.symbol("COLUMN_PROPERTY")  # type: ignore
+SYNONYM_PROPERTY: int = util.symbol("SYNONYM_PROPERTY")  # type: ignore
+COMPOSITE_PROPERTY: int = util.symbol("COMPOSITE_PROPERTY")  # type: ignore
+DECLARED_ATTR: int = util.symbol("DECLARED_ATTR")  # type: ignore
+MAPPER_PROPERTY: int = util.symbol("MAPPER_PROPERTY")  # type: ignore
+AS_DECLARATIVE: int = util.symbol("AS_DECLARATIVE")  # type: ignore
+AS_DECLARATIVE_BASE: int = util.symbol("AS_DECLARATIVE_BASE")  # type: ignore
+DECLARATIVE_MIXIN: int = util.symbol("DECLARATIVE_MIXIN")  # type: ignore
+
+_lookup: Dict[str, Tuple[int, Set[str]]] = {
     "Column": (
         COLUMN,
         {
@@ -145,7 +153,21 @@ _lookup = {
 }
 
 
-def _mro_has_id(mro: List[TypeInfo], type_id: int):
+def _has_base_type_id(info: TypeInfo, type_id: int) -> bool:
+    for mr in info.mro:
+        check_type_id, fullnames = _lookup.get(mr.name, (None, None))
+        if check_type_id == type_id:
+            break
+    else:
+        return False
+
+    if fullnames is None:
+        return False
+
+    return mr.fullname in fullnames
+
+
+def _mro_has_id(mro: List[TypeInfo], type_id: int) -> bool:
     for mr in mro:
         check_type_id, fullnames = _lookup.get(mr.name, (None, None))
         if check_type_id == type_id:
@@ -153,65 +175,75 @@ def _mro_has_id(mro: List[TypeInfo], type_id: int):
     else:
         return False
 
+    if fullnames is None:
+        return False
+
     return mr.fullname in fullnames
 
 
 def _type_id_for_unbound_type(
     type_: UnboundType, cls: ClassDef, api: SemanticAnalyzerPluginInterface
-) -> int:
+) -> Optional[int]:
     type_id = None
 
     sym = api.lookup_qualified(type_.name, type_)
     if sym is not None:
         if isinstance(sym.node, TypeAlias):
-            type_id = _type_id_for_named_node(sym.node.target.type)
+            target_type = get_proper_type(sym.node.target)
+            if isinstance(target_type, Instance):
+                type_id = _type_id_for_named_node(target_type.type)
         elif isinstance(sym.node, TypeInfo):
             type_id = _type_id_for_named_node(sym.node)
 
     return type_id
 
 
-def _type_id_for_callee(callee: Expression) -> int:
-    if isinstance(callee.node, FuncDef):
-        return _type_id_for_funcdef(callee.node)
-    elif isinstance(callee.node, TypeAlias):
-        type_id = _type_id_for_fullname(callee.node.target.type.fullname)
-    elif isinstance(callee.node, TypeInfo):
-        type_id = _type_id_for_named_node(callee)
-    else:
-        type_id = None
+def _type_id_for_callee(callee: Expression) -> Optional[int]:
+    if isinstance(callee, (MemberExpr, NameExpr)):
+        if isinstance(callee.node, FuncDef):
+            return _type_id_for_funcdef(callee.node)
+        elif isinstance(callee.node, TypeAlias):
+            target_type = get_proper_type(callee.node.target)
+            if isinstance(target_type, Instance):
+                type_id = _type_id_for_fullname(target_type.type.fullname)
+        elif isinstance(callee.node, TypeInfo):
+            type_id = _type_id_for_named_node(callee)
+        else:
+            type_id = None
     return type_id
 
 
-def _type_id_for_funcdef(node: FuncDef) -> int:
-    if hasattr(node.type.ret_type, "type"):
-        type_id = _type_id_for_fullname(node.type.ret_type.type.fullname)
-    else:
-        type_id = None
-    return type_id
+def _type_id_for_funcdef(node: FuncDef) -> Optional[int]:
+    if node.type and isinstance(node.type, CallableType):
+        ret_type = get_proper_type(node.type.ret_type)
+
+        if isinstance(ret_type, Instance):
+            return _type_id_for_fullname(ret_type.type.fullname)
+
+    return None
 
 
-def _type_id_for_named_node(node: Union[RefExpr, SymbolNode]) -> int:
+def _type_id_for_named_node(
+    node: Union[NameExpr, MemberExpr, SymbolNode]
+) -> Optional[int]:
     type_id, fullnames = _lookup.get(node.name, (None, None))
 
-    if type_id is None:
+    if type_id is None or fullnames is None:
         return None
-
     elif node.fullname in fullnames:
         return type_id
     else:
         return None
 
 
-def _type_id_for_fullname(fullname: str) -> int:
+def _type_id_for_fullname(fullname: str) -> Optional[int]:
     tokens = fullname.split(".")
     immediate = tokens[-1]
 
     type_id, fullnames = _lookup.get(immediate, (None, None))
 
-    if type_id is None:
+    if type_id is None or fullnames is None:
         return None
-
     elif fullname in fullnames:
         return type_id
     else: