Cache package names in create_source_list (#4848)

msullivan · web-flow · commit 6b05cb7d6029 · 2018-04-04T11:16:08.000-07:00
diff --git a/mypy/find_sources.py b/mypy/find_sources.py
@@ -2,7 +2,7 @@
 
 import os.path
 
-from typing import List, Sequence, Set, Tuple, Optional
+from typing import List, Sequence, Set, Tuple, Optional, Dict
 
 from mypy.build import BuildSource, PYTHON_EXTENSIONS
 from mypy.fscache import FileSystemMetaCache
@@ -17,20 +17,23 @@ class InvalidSourceList(Exception):
 
 
 def create_source_list(files: Sequence[str], options: Options,
-                       fscache: Optional[FileSystemMetaCache] = None) -> List[BuildSource]:
+                       fscache: Optional[FileSystemMetaCache] = None,
+                       allow_empty_dir: bool = False) -> List[BuildSource]:
     """From a list of source files/directories, makes a list of BuildSources.
 
     Raises InvalidSourceList on errors.
     """
     fscache = fscache or FileSystemMetaCache()
+    finder = SourceFinder(fscache)
+
     targets = []
     for f in files:
         if f.endswith(PY_EXTENSIONS):
             # Can raise InvalidSourceList if a directory doesn't have a valid module name.
-            targets.append(BuildSource(f, crawl_up(fscache, f)[1], None))
+            targets.append(BuildSource(f, finder.crawl_up(f), None))
         elif fscache.isdir(f):
-            sub_targets = expand_dir(fscache, f)
-            if not sub_targets:
+            sub_targets = finder.expand_dir(f)
+            if not sub_targets and not allow_empty_dir:
                 raise InvalidSourceList("There are no .py[i] files in directory '{}'"
                                         .format(f))
             targets.extend(sub_targets)
@@ -52,60 +55,101 @@ def keyfunc(name: str) -> Tuple[int, str]:
     return (-1, name)
 
 
-def expand_dir(fscache: FileSystemMetaCache,
-               arg: str, mod_prefix: str = '') -> List[BuildSource]:
-    """Convert a directory name to a list of sources to build."""
-    f = get_init_file(fscache, arg)
-    if mod_prefix and not f:
-        return []
-    seen = set()  # type: Set[str]
-    sources = []
-    if f and not mod_prefix:
-        top_dir, top_mod = crawl_up(fscache, f)
-        mod_prefix = top_mod + '.'
-    if mod_prefix:
-        sources.append(BuildSource(f, mod_prefix.rstrip('.'), None))
-    names = fscache.listdir(arg)
-    names.sort(key=keyfunc)
-    for name in names:
-        path = os.path.join(arg, name)
-        if fscache.isdir(path):
-            sub_sources = expand_dir(fscache, path, mod_prefix + name + '.')
-            if sub_sources:
-                seen.add(name)
-                sources.extend(sub_sources)
+class SourceFinder:
+    def __init__(self, fscache: FileSystemMetaCache) -> None:
+        self.fscache = fscache
+        # A cache for package names, mapping from module id to directory path
+        self.package_cache = {}  # type: Dict[str, str]
+
+    def expand_dir(self, arg: str, mod_prefix: str = '') -> List[BuildSource]:
+        """Convert a directory name to a list of sources to build."""
+        f = self.get_init_file(arg)
+        if mod_prefix and not f:
+            return []
+        seen = set()  # type: Set[str]
+        sources = []
+        if f and not mod_prefix:
+            top_mod = self.crawl_up(f)
+            mod_prefix = top_mod + '.'
+        if mod_prefix:
+            sources.append(BuildSource(f, mod_prefix.rstrip('.'), None))
+        names = self.fscache.listdir(arg)
+        names.sort(key=keyfunc)
+        for name in names:
+            path = os.path.join(arg, name)
+            if self.fscache.isdir(path):
+                sub_sources = self.expand_dir(path, mod_prefix + name + '.')
+                if sub_sources:
+                    seen.add(name)
+                    sources.extend(sub_sources)
+            else:
+                base, suffix = os.path.splitext(name)
+                if base == '__init__':
+                    continue
+                if base not in seen and '.' not in base and suffix in PY_EXTENSIONS:
+                    seen.add(base)
+                    src = BuildSource(path, mod_prefix + base, None)
+                    sources.append(src)
+        return sources
+
+    def crawl_up(self, arg: str) -> str:
+        """Given a .py[i] filename, return module.
+
+        We crawl up the path until we find a directory without
+        __init__.py[i], or until we run out of path components.
+        """
+        dir, mod = os.path.split(arg)
+        mod = strip_py(mod) or mod
+        base = self.crawl_up_dir(dir)
+        if mod == '__init__' or not mod:
+            mod = base
         else:
-            base, suffix = os.path.splitext(name)
-            if base == '__init__':
-                continue
-            if base not in seen and '.' not in base and suffix in PY_EXTENSIONS:
-                seen.add(base)
-                src = BuildSource(path, mod_prefix + base, None)
-                sources.append(src)
-    return sources
+            mod = module_join(base, mod)
 
+        return mod
 
-def crawl_up(fscache: FileSystemMetaCache, arg: str) -> Tuple[str, str]:
-    """Given a .py[i] filename, return (root directory, module).
+    def crawl_up_dir(self, dir: str) -> str:
+        """Given a directory name, return the corresponding module name.
 
-    We crawl up the path until we find a directory without
-    __init__.py[i], or until we run out of path components.
-    """
-    dir, mod = os.path.split(arg)
-    mod = strip_py(mod) or mod
-    while dir and get_init_file(fscache, dir):
-        dir, base = os.path.split(dir)
-        if not base:
-            break
-        # Ensure that base is a valid python module name
-        if not base.isidentifier():
-            raise InvalidSourceList('{} is not a valid Python package name'.format(base))
-        if mod == '__init__' or not mod:
-            mod = base
+        Use package_cache to cache results.
+        """
+        if dir in self.package_cache:
+            return self.package_cache[dir]
+
+        parent_dir, base = os.path.split(dir)
+        if not dir or not self.get_init_file(dir) or not base:
+            res = ''
         else:
-            mod = base + '.' + mod
+            # Ensure that base is a valid python module name
+            if not base.isidentifier():
+                raise InvalidSourceList('{} is not a valid Python package name'.format(base))
+            parent = self.crawl_up_dir(parent_dir)
+            res = module_join(parent, base)
+
+        self.package_cache[dir] = res
+        return res
+
+    def get_init_file(self, dir: str) -> Optional[str]:
+        """Check whether a directory contains a file named __init__.py[i].
+
+        If so, return the file's name (with dir prefixed).  If not, return
+        None.
 
-    return dir, mod
+        This prefers .pyi over .py (because of the ordering of PY_EXTENSIONS).
+        """
+        for ext in PY_EXTENSIONS:
+            f = os.path.join(dir, '__init__' + ext)
+            if self.fscache.isfile(f):
+                return f
+        return None
+
+
+def module_join(parent: str, child: str) -> str:
+    """Join module ids, accounting for a possibly empty parent."""
+    if parent:
+        return parent + '.' + child
+    else:
+        return child
 
 
 def strip_py(arg: str) -> Optional[str]:
@@ -117,18 +161,3 @@ def strip_py(arg: str) -> Optional[str]:
         if arg.endswith(ext):
             return arg[:-len(ext)]
     return None
-
-
-def get_init_file(fscache: FileSystemMetaCache, dir: str) -> Optional[str]:
-    """Check whether a directory contains a file named __init__.py[i].
-
-    If so, return the file's name (with dir prefixed).  If not, return
-    None.
-
-    This prefers .pyi over .py (because of the ordering of PY_EXTENSIONS).
-    """
-    for ext in PY_EXTENSIONS:
-        f = os.path.join(dir, '__init__' + ext)
-        if fscache.isfile(f):
-            return f
-    return None
diff --git a/mypy/test/testfinegrained.py b/mypy/test/testfinegrained.py
@@ -27,7 +27,7 @@
 from mypy.server.mergecheck import check_consistency
 from mypy.dmypy_server import Server
 from mypy.main import parse_config_file
-from mypy.find_sources import expand_dir, create_source_list
+from mypy.find_sources import create_source_list
 from mypy.fscache import FileSystemMetaCache
 
 import pytest  # type: ignore  # no pytest in typeshed
@@ -238,7 +238,8 @@ def parse_sources(self, program_text: str,
             base = BuildSource(os.path.join(test_temp_dir, 'main'), '__main__', None)
             # Use expand_dir instead of create_source_list to avoid complaints
             # when there aren't any .py files in an increment
-            return [base] + expand_dir(FileSystemMetaCache(), test_temp_dir)
+            return [base] + create_source_list([test_temp_dir], options,
+                                               allow_empty_dir=True)
 
 
 def normalize_messages(messages: List[str]) -> List[str]:
