Work towards following imports in mypy daemon (#8559)

This implements some basic use cases, such as adding files, removing files,
and basic package support.

The implementation is separate from what is used for coarse-grained
incremental mode, since it seemed too difficult (and risky) to try to
merge them.

Many special cases are not handled correctly yet, and code is in need of some
cleanup. Some functionality is still missing.

Following imports is only enabled in tests for now. I'll enable the feature
generally after it's more fully implemented.

Author: JukkaL

mypy/dmypy_server.py

@@ -16,19 +16,19 @@
 import traceback
 from contextlib import redirect_stderr, redirect_stdout
 
-from typing import AbstractSet, Any, Callable, Dict, List, Optional, Sequence, Tuple
+from typing import AbstractSet, Any, Callable, Dict, List, Optional, Sequence, Tuple, Set
 from typing_extensions import Final
 
 import mypy.build
 import mypy.errors
 import mypy.main
 from mypy.find_sources import create_source_list, InvalidSourceList
-from mypy.server.update import FineGrainedBuildManager
+from mypy.server.update import FineGrainedBuildManager, refresh_suppressed_submodules
 from mypy.dmypy_util import receive
 from mypy.ipc import IPCServer
 from mypy.fscache import FileSystemCache
 from mypy.fswatcher import FileSystemWatcher, FileData
-from mypy.modulefinder import BuildSource, compute_search_paths
+from mypy.modulefinder import BuildSource, compute_search_paths, FindModuleCache, SearchPaths
 from mypy.options import Options
 from mypy.suggestions import SuggestionFailure, SuggestionEngine
 from mypy.typestate import reset_global_state
@@ -361,8 +361,12 @@ def cmd_recheck(self,
         t1 = time.time()
         manager = self.fine_grained_manager.manager
         manager.log("fine-grained increment: cmd_recheck: {:.3f}s".format(t1 - t0))
-        res = self.fine_grained_increment(sources, is_tty, terminal_width,
-                                          remove, update)
+        if not self.following_imports():
+            messages = self.fine_grained_increment(sources, remove, update)
+        else:
+            assert remove is None and update is None
+            messages = self.fine_grained_increment_follow_imports(sources)
+        res = self.increment_output(messages, sources, is_tty, terminal_width)
         self.fscache.flush()
         self.update_stats(res)
         return res
@@ -377,7 +381,11 @@ def check(self, sources: List[BuildSource],
         if not self.fine_grained_manager:
             res = self.initialize_fine_grained(sources, is_tty, terminal_width)
         else:
-            res = self.fine_grained_increment(sources, is_tty, terminal_width)
+            if not self.following_imports():
+                messages = self.fine_grained_increment(sources)
+            else:
+                messages = self.fine_grained_increment_follow_imports(sources)
+            res = self.increment_output(messages, sources, is_tty, terminal_width)
         self.fscache.flush()
         self.update_stats(res)
         return res
@@ -389,6 +397,11 @@ def update_stats(self, res: Dict[str, Any]) -> None:
             res['stats'] = manager.stats
             manager.stats = {}
 
+    def following_imports(self) -> bool:
+        """Are we following imports?"""
+        # TODO: What about silent?
+        return self.options.follow_imports == 'normal'
+
     def initialize_fine_grained(self, sources: List[BuildSource],
                                 is_tty: bool, terminal_width: int) -> Dict[str, Any]:
         self.fswatcher = FileSystemWatcher(self.fscache)
@@ -408,6 +421,11 @@ def initialize_fine_grained(self, sources: List[BuildSource],
             return {'out': out, 'err': err, 'status': 2}
         messages = result.errors
         self.fine_grained_manager = FineGrainedBuildManager(result)
+
+        if self.following_imports():
+            sources = find_all_sources_in_build(self.fine_grained_manager.graph)
+            self.update_sources(sources)
+
         self.previous_sources = sources
 
         # If we are using the fine-grained cache, build hasn't actually done
@@ -436,13 +454,19 @@ def initialize_fine_grained(self, sources: List[BuildSource],
             t3 = time.time()
             # Run an update
             messages = self.fine_grained_manager.update(changed, removed)
+
+            if self.following_imports():
+                # We need to do another update to any new files found by following imports.
+                messages = self.fine_grained_increment_follow_imports(sources)
+
             t4 = time.time()
             self.fine_grained_manager.manager.add_stats(
                 update_sources_time=t1 - t0,
                 build_time=t2 - t1,
                 find_changes_time=t3 - t2,
                 fg_update_time=t4 - t3,
                 files_changed=len(removed) + len(changed))
+
         else:
             # Stores the initial state of sources as a side effect.
             self.fswatcher.find_changed()
@@ -457,11 +481,19 @@ def initialize_fine_grained(self, sources: List[BuildSource],
 
     def fine_grained_increment(self,
                                sources: List[BuildSource],
-                               is_tty: bool,
-                               terminal_width: int,
                                remove: Optional[List[str]] = None,
                                update: Optional[List[str]] = None,
-                               ) -> Dict[str, Any]:
+                               ) -> List[str]:
+        """Perform a fine-grained type checking increment.
+
+        If remove and update are None, determine changed paths by using
+        fswatcher. Otherwise, assume that only these files have changes.
+
+        Args:
+            sources: sources passed on the command line
+            remove: paths of files that have been removed
+            update: paths of files that have been changed or created
+        """
         assert self.fine_grained_manager is not None
         manager = self.fine_grained_manager.manager
 
@@ -486,8 +518,180 @@ def fine_grained_increment(self,
             fg_update_time=t2 - t1,
             files_changed=len(removed) + len(changed))
 
-        status = 1 if messages else 0
         self.previous_sources = sources
+        return messages
+
+    def fine_grained_increment_follow_imports(self, sources: List[BuildSource]) -> List[str]:
+        """Like fine_grained_increment, but follow imports."""
+
+        # TODO:
+        #  - file events
+        #  - search path updates
+        #  - logging
+
+        changed_paths = self.fswatcher.find_changed()
+
+        assert self.fine_grained_manager is not None
+        fine_grained_manager = self.fine_grained_manager
+        graph = fine_grained_manager.graph
+        manager = fine_grained_manager.manager
+
+        orig_modules = list(graph.keys())
+
+        # TODO: Are the differences from fine_grained_increment(), such as
+        #       updating sources after finding changed, necessary?
+        self.update_sources(sources)
+
+        sources_set = {source.module for source in sources}
+
+        # Find changed modules reachable from roots (or in roots) already in graph.
+        seen = set()  # type: Set[str]
+        changed, removed, new_files = self.follow_imports(
+            sources, graph, seen, changed_paths, sources_set
+        )
+        sources.extend(new_files)
+
+        # Process changes directly reachable from roots.
+        messages = fine_grained_manager.update(changed, removed)
+
+        # Follow deps from changed modules (still within graph).
+        worklist = changed[:]
+        while worklist:
+            module = worklist.pop()
+            if module[0] not in graph:
+                continue
+            sources2 = self.direct_imports(module, graph)
+            changed, removed, new_files = self.follow_imports(
+                sources2, graph, seen, changed_paths, sources_set
+            )
+            sources.extend(new_files)
+            self.update_sources(new_files)
+            messages = fine_grained_manager.update(changed, removed)
+            # TODO: Removed?
+            worklist.extend(changed)
+
+        for module_id, state in graph.items():
+            refresh_suppressed_submodules(module_id, state.path, fine_grained_manager.deps, graph,
+                                          self.fscache)
+
+        # There may be new files that became available, currently treated as
+        # suppressed imports. Process them.
+        seen_suppressed = set()  # type: Set[str]
+        while True:
+            # TODO: Merge seen and seen_suppressed?
+            new_unsuppressed, seen_suppressed = self.find_added_suppressed(
+                graph, seen_suppressed, manager.search_paths
+            )
+            if not new_unsuppressed:
+                break
+            new_files = [BuildSource(mod[1], mod[0]) for mod in new_unsuppressed]
+            sources.extend(new_files)
+            self.update_sources(new_files)
+            messages = fine_grained_manager.update(new_unsuppressed, [])
+
+            for module_id, path in new_unsuppressed:
+                refresh_suppressed_submodules(module_id, path, fine_grained_manager.deps, graph,
+                                              self.fscache)
+
+        # Find all original modules in graph that were not reached -- they are deleted.
+        to_delete = []
+        for module_id in orig_modules:
+            if module_id not in graph:
+                continue
+            if module_id not in seen and module_id not in seen_suppressed:
+                module_path = graph[module_id].path
+                assert module_path is not None
+                to_delete.append((module_id, module_path))
+        if to_delete:
+            messages = fine_grained_manager.update([], to_delete)
+
+        fix_module_deps(graph)
+
+        # Store current file state as side effect
+        self.fswatcher.find_changed()
+
+        self.previous_sources = find_all_sources_in_build(graph)
+        self.update_sources(self.previous_sources)
+        return messages
+
+    def follow_imports(self,
+                       sources: List[BuildSource],
+                       graph: mypy.build.Graph,
+                       seen: Set[str],
+                       changed_paths: AbstractSet[str],
+                       sources_set: Set[str]) -> Tuple[List[Tuple[str, str]],
+                                                       List[Tuple[str, str]],
+                                                       List[BuildSource]]:
+        """Follow imports within graph from given sources.
+
+        Args:
+            sources: roots of modules to search
+            graph: module graph to use for the search
+            seen: modules we've seen before that won't be visited (mutated here!)
+            changed_paths: which paths have changed (stop search here and return any found)
+            sources_set: set of sources (TODO: relationship with seen)
+
+        Return (reachable changed modules, removed modules, updated file list).
+        """
+        changed = []
+        new_files = []
+        worklist = sources[:]
+        seen.update(source.module for source in worklist)
+        while worklist:
+            nxt = worklist.pop()
+            if nxt.module not in sources_set:
+                sources_set.add(nxt.module)
+                new_files.append(nxt)
+            if nxt.path in changed_paths:
+                assert nxt.path is not None  # TODO
+                changed.append((nxt.module, nxt.path))
+            elif nxt.module in graph:
+                state = graph[nxt.module]
+                for dep in state.dependencies:
+                    if dep not in seen:
+                        seen.add(dep)
+                        worklist.append(BuildSource(graph[dep].path,
+                                                    graph[dep].id))
+        return changed, [], new_files
+
+    def direct_imports(self,
+                       module: Tuple[str, str],
+                       graph: mypy.build.Graph) -> List[BuildSource]:
+        """Return the direct imports of module not included in seen."""
+        state = graph[module[0]]
+        return [BuildSource(graph[dep].path, dep)
+                for dep in state.dependencies]
+
+    def find_added_suppressed(self,
+                              graph: mypy.build.Graph,
+                              seen: Set[str],
+                              search_paths: SearchPaths) -> Tuple[List[Tuple[str, str]], Set[str]]:
+        """Find suppressed modules that have been added (and not included in seen)."""
+        all_suppressed = set()
+        for module, state in graph.items():
+            all_suppressed |= state.suppressed_set
+
+        # TODO: Namespace packages
+        # TODO: Handle seen?
+
+        finder = FindModuleCache(search_paths, self.fscache, self.options)
+
+        found = []
+
+        for module in all_suppressed:
+            result = finder.find_module(module)
+            if isinstance(result, str) and module not in seen:
+                found.append((module, result))
+                seen.add(module)
+
+        return found, seen
+
+    def increment_output(self,
+                         messages: List[str],
+                         sources: List[BuildSource],
+                         is_tty: bool,
+                         terminal_width: int) -> Dict[str, Any]:
+        status = 1 if messages else 0
         messages = self.pretty_messages(messages, len(sources), is_tty, terminal_width)
         return {'out': ''.join(s + '\n' for s in messages), 'err': '', 'status': status}
 
@@ -622,3 +826,30 @@ def get_meminfo() -> Dict[str, Any]:
                 factor = 1024  # Linux
             res['memory_maxrss_mib'] = rusage.ru_maxrss * factor / MiB
     return res
+
+
+def find_all_sources_in_build(graph: mypy.build.Graph) -> List[BuildSource]:
+    result = []
+    for module, state in graph.items():
+        result.append(BuildSource(state.path, module))
+    return result
+
+
+def fix_module_deps(graph: mypy.build.Graph) -> None:
+    """After an incremental update, update module dependencies to reflect the new state.
+
+    This can make some suppressed dependencies non-suppressed, and vice versa (if modules
+    have been added to or removed from the build).
+    """
+    for module, state in graph.items():
+        new_suppressed = []
+        new_dependencies = []
+        for dep in state.dependencies + state.suppressed:
+            if dep in graph:
+                new_dependencies.append(dep)
+            else:
+                new_suppressed.append(dep)
+        state.dependencies = new_dependencies
+        state.dependencies_set = set(new_dependencies)
+        state.suppressed = new_suppressed
+        state.suppressed_set = set(new_suppressed)

mypy/modulefinder.py

@@ -70,7 +70,7 @@ class BuildSource:
     """A single source file."""
 
     def __init__(self, path: Optional[str], module: Optional[str],
-                 text: Optional[str], base_dir: Optional[str] = None) -> None:
+                 text: Optional[str] = None, base_dir: Optional[str] = None) -> None:
         self.path = path  # File where it's found (e.g. 'xxx/yyy/foo/bar.py')
         self.module = module or '__main__'  # Module name (e.g. 'foo.bar')
         self.text = text  # Source code, if initially supplied, else None