applytype: use zip, closure

python · JukkaL · Sep 8, 2020 · Apr 7, 2020 · Apr 7, 2020 · Apr 7, 2020
commit fe7b65547d4a1b76b367e22f84542f32b85d6e6c
diff --git a/mypy/applytype.py b/mypy/applytype.py
@@ -4,7 +4,8 @@
 import mypy.sametypes
 from mypy.expandtype import expand_type
 from mypy.types import (
-    Type, TypeVarId, TypeVarType, CallableType, AnyType, PartialType, get_proper_types
+    Type, TypeVarId, TypeVarType, CallableType, AnyType, PartialType, get_proper_types,
+    TypeVarDef, TypeVarLikeDef, ProperType
 )
 from mypy.nodes import Context
 
@@ -29,20 +30,23 @@ def apply_generic_arguments(
     # Check that inferred type variable values are compatible with allowed
     # values and bounds.  Also, promote subtype values to allowed values.
     types = get_proper_types(orig_types)
-    for i, type in enumerate(types):
-        assert not isinstance(type, PartialType), "Internal error: must never apply partial type"
-        values = get_proper_types(callable.variables[i].values)
-        if type is None:
-            continue
+
+    # Create a map from type variable id to target type.
+    id_to_type = {}  # type: Dict[TypeVarId, Type]
+
+    def get_target_type(tvar: TypeVarLikeDef, type: ProperType) -> Optional[Type]:
+        # TODO(shantanu): fix for ParamSpecDef
+        assert isinstance(tvar, TypeVarDef)
+        values = get_proper_types(tvar.values)
         if values:
             if isinstance(type, AnyType):
-                continue
+                return type
             if isinstance(type, TypeVarType) and type.values:
                 # Allow substituting T1 for T if every allowed value of T1
                 # is also a legal value of T.
                 if all(any(mypy.sametypes.is_same_type(v, v1) for v in values)
                        for v1 in type.values):
-                    continue
+                    return type
             matching = []
             for value in values:
                 if mypy.subtypes.is_subtype(type, value):
@@ -53,28 +57,26 @@ def apply_generic_arguments(
                 for match in matching[1:]:
                     if mypy.subtypes.is_subtype(match, best):
                         best = match
-                types[i] = best
-            else:
-                if skip_unsatisfied:
-                    types[i] = None
-                else:
-                    report_incompatible_typevar_value(callable, type, callable.variables[i].name,
-                                                      context)
+                return best
+            if skip_unsatisfied:
+                return None
+            report_incompatible_typevar_value(callable, type, tvar.name, context)
         else:
-            upper_bound = callable.variables[i].upper_bound
+            upper_bound = tvar.upper_bound
             if not mypy.subtypes.is_subtype(type, upper_bound):
                 if skip_unsatisfied:
-                    types[i] = None
-                else:
-                    report_incompatible_typevar_value(callable, type, callable.variables[i].name,
-                                                      context)
+                    return None
+                report_incompatible_typevar_value(callable, type, tvar.name, context)
+        return type
 
-    # Create a map from type variable id to target type.
-    id_to_type = {}  # type: Dict[TypeVarId, Type]
-    for i, tv in enumerate(tvars):
-        typ = types[i]
-        if typ:
-            id_to_type[tv.id] = typ
+    for tvar, type in zip(tvars, types):
+        assert not isinstance(type, PartialType), "Internal error: must never apply partial type"
+        if type is None:
+            continue
+
+        target_type = get_target_type(tvar, type)
+        if target_type is not None:
+            id_to_type[tvar.id] = target_type
 
     # Apply arguments to argument types.
     arg_types = [expand_type(at, id_to_type) for at in callable.arg_types]