Add a new, lower priority for imports inside "if MYPY"

Author: Guido van Rossum

mypy/build.py

@@ -24,7 +24,8 @@
 from typing import (AbstractSet, Dict, Iterable, Iterator, List,
                     NamedTuple, Optional, Set, Tuple, Union)
 
-from mypy.nodes import (MypyFile, Import, ImportFrom, ImportAll)
+from mypy.types import Type
+from mypy.nodes import (MypyFile, Node, ImportBase, Import, ImportFrom, ImportAll)
 from mypy.semanal import FirstPass, SemanticAnalyzer, ThirdPass
 from mypy.checker import TypeChecker
 from mypy.indirection import TypeIndirectionVisitor
@@ -307,9 +308,22 @@ def default_lib_path(data_dir: str, pyversion: Tuple[int, int]) -> List[str]:
 PRI_MED = 10  # top-level "import X"
 PRI_LOW = 20  # either form inside a function
 PRI_INDIRECT = 30  # an indirect dependency
+PRI_MYPY = 40  # inside "if MYPY" or "if typing.TYPE_CHECKING"
 PRI_ALL = 99  # include all priorities
 
 
+def import_priority(imp: ImportBase, toplevel_priority: int) -> int:
+    """Compute import priority from an import node."""
+    if imp.is_mypy_only:
+        # Inside "if MYPY" or "if typing.TYPE_CHECKING"
+        return PRI_MYPY
+    if not imp.is_top_level:
+        # Inside a function
+        return PRI_LOW
+    # A regular import; priority determined by argument.
+    return toplevel_priority
+
+
 # TODO: Get rid of all_types.  It's not used except for one log message.
 #       Maybe we could instead publish a map from module ID to its type_map.
 class BuildManager:
@@ -395,20 +409,21 @@ def correct_rel_imp(imp: Union[ImportFrom, ImportAll]) -> str:
         for imp in file.imports:
             if not imp.is_unreachable:
                 if isinstance(imp, Import):
-                    pri = PRI_MED if imp.is_top_level else PRI_LOW
+                    pri = import_priority(imp, PRI_MED)
+                    ancestor_pri = import_priority(imp, PRI_LOW)
                     for id, _ in imp.ids:
                         ancestor_parts = id.split(".")[:-1]
                         ancestors = []
                         for part in ancestor_parts:
                             ancestors.append(part)
-                            res.append((PRI_LOW, ".".join(ancestors), imp.line))
+                            res.append((ancestor_pri, ".".join(ancestors), imp.line))
                         res.append((pri, id, imp.line))
                 elif isinstance(imp, ImportFrom):
                     cur_id = correct_rel_imp(imp)
                     pos = len(res)
                     all_are_submodules = True
                     # Also add any imported names that are submodules.
-                    pri = PRI_MED if imp.is_top_level else PRI_LOW
+                    pri = import_priority(imp, PRI_MED)
                     for name, __ in imp.names:
                         sub_id = cur_id + '.' + name
                         if self.is_module(sub_id):
@@ -421,10 +436,10 @@ def correct_rel_imp(imp: Union[ImportFrom, ImportAll]) -> str:
                     # cur_id is also a dependency, and we should
                     # insert it *before* any submodules.
                     if not all_are_submodules:
-                        pri = PRI_HIGH if imp.is_top_level else PRI_LOW
+                        pri = import_priority(imp, PRI_HIGH)
                         res.insert(pos, ((pri, cur_id, imp.line)))
                 elif isinstance(imp, ImportAll):
-                    pri = PRI_HIGH if imp.is_top_level else PRI_LOW
+                    pri = import_priority(imp, PRI_HIGH)
                     res.append((pri, correct_rel_imp(imp), imp.line))
 
         return res

mypy/nodes.py

@@ -285,8 +285,11 @@ def deserialize(cls, data: JsonDict) -> 'MypyFile':
 
 class ImportBase(Statement):
     """Base class for all import statements."""
-    is_unreachable = False
-    is_top_level = False  # Set by semanal.FirstPass
+
+    is_unreachable = False  # Set by semanal.FirstPass if inside `if False` etc.
+    is_top_level = False  # Ditto if outside any class or def
+    is_mypy_only = False  # Ditto if inside `if TYPE_CHECKING` or `if MYPY`
+
     # If an import replaces existing definitions, we construct dummy assignment
     # statements that assign the imported names to the names in the current scope,
     # for type checking purposes. Example:

mypy/semanal.py

@@ -84,10 +84,20 @@
 T = TypeVar('T')
 
 
-# Inferred value of an expression.
-ALWAYS_TRUE = 0
-ALWAYS_FALSE = 1
-TRUTH_VALUE_UNKNOWN = 2
+# Inferred truth value of an expression.
+ALWAYS_TRUE = 1
+MYPY_TRUE = 2  # True in mypy, False at runtime
+ALWAYS_FALSE = 3
+MYPY_FALSE = 4  # False in mypy, True at runtime
+TRUTH_VALUE_UNKNOWN = 5
+
+inverted_truth_mapping = {
+    ALWAYS_TRUE: ALWAYS_FALSE,
+    ALWAYS_FALSE: ALWAYS_TRUE,
+    TRUTH_VALUE_UNKNOWN: TRUTH_VALUE_UNKNOWN,
+    MYPY_TRUE: MYPY_FALSE,
+    MYPY_FALSE: MYPY_TRUE,
+}
 
 # Map from obsolete name to the current spelling.
 obsolete_name_mapping = {
@@ -3082,12 +3092,16 @@ def infer_reachability_of_if_statement(s: IfStmt,
                                        platform: str) -> None:
     for i in range(len(s.expr)):
         result = infer_if_condition_value(s.expr[i], pyversion, platform)
-        if result == ALWAYS_FALSE:
+        if result in (ALWAYS_FALSE, MYPY_FALSE):
             # The condition is always false, so we skip the if/elif body.
             mark_block_unreachable(s.body[i])
-        elif result == ALWAYS_TRUE:
+        elif result in (ALWAYS_TRUE, MYPY_TRUE):
             # This condition is always true, so all of the remaining
             # elif/else bodies will never be executed.
+            if result == MYPY_TRUE:
+                # This condition is false at runtime; this will affect
+                # import priorities.
+                mark_block_mypy_only(s.body[i])
             for body in s.body[i + 1:]:
                 mark_block_unreachable(body)
             if s.else_body:
@@ -3099,7 +3113,8 @@ def infer_if_condition_value(expr: Expression, pyversion: Tuple[int, int], platf
     """Infer whether if condition is always true/false.
 
     Return ALWAYS_TRUE if always true, ALWAYS_FALSE if always false,
-    and TRUTH_VALUE_UNKNOWN otherwise.
+    MYPY_TRUE if true under mypy and false at runtime, MYPY_FALSE if
+    false under mypy and true at runtime, else TRUTH_VALUE_UNKNOWN.
     """
     name = ''
     negated = False
@@ -3123,12 +3138,9 @@ def infer_if_condition_value(expr: Expression, pyversion: Tuple[int, int], platf
         elif name == 'PY3':
             result = ALWAYS_TRUE if pyversion[0] == 3 else ALWAYS_FALSE
         elif name == 'MYPY' or name == 'TYPE_CHECKING':
-            result = ALWAYS_TRUE
+            result = MYPY_TRUE
     if negated:
-        if result == ALWAYS_TRUE:
-            result = ALWAYS_FALSE
-        elif result == ALWAYS_FALSE:
-            result = ALWAYS_TRUE
+        result = inverted_truth_mapping[result]
     return result
 
 
@@ -3303,6 +3315,23 @@ def visit_import_all(self, node: ImportAll) -> None:
         node.is_unreachable = True
 
 
+def mark_block_mypy_only(block: Block) -> None:
+    block.accept(MarkImportsMypyOnlyVisitor())
+
+
+class MarkImportsMypyOnlyVisitor(TraverserVisitor):
+    """Visitor that sets is_mypy_only (which affects priority)."""
+
+    def visit_import(self, node: Import) -> None:
+        node.is_mypy_only = True
+
+    def visit_import_from(self, node: ImportFrom) -> None:
+        node.is_mypy_only = True
+
+    def visit_import_all(self, node: ImportAll) -> None:
+        node.is_mypy_only = True
+
+
 def is_identity_signature(sig: Type) -> bool:
     """Is type a callable of form T -> T (where T is a type variable)?"""
     if isinstance(sig, CallableType) and sig.arg_kinds == [ARG_POS]: