First implementation

Signed-off-by: Bernat Gabor <bgabor8@bloomberg.net>

mypy/build.py

@@ -52,6 +52,7 @@
 from mypy.options import Options
 from mypy.parse import parse
 from mypy.stats import dump_type_stats
+from mypy.stub_src_merge import MergeFiles
 from mypy.types import Type
 from mypy.version import __version__
 from mypy.plugin import Plugin, ChainedPlugin, plugin_types
@@ -107,6 +108,8 @@ def __init__(self, sources: List[BuildSource]) -> None:
                 self.source_text_present = True
             elif source.path:
                 self.source_paths.add(source.path)
+                if source.merge_with and source.merge_with.path:
+                    self.source_paths.add(source.merge_with.path)
             else:
                 self.source_modules.add(source.module)
 
@@ -2154,6 +2157,23 @@ def generate_unused_ignore_notes(self) -> None:
                 self.verify_dependencies(suppressed_only=True)
             self.manager.errors.generate_unused_ignore_notes(self.xpath)
 
+    def merge_with(self, state: 'State', errors: Errors) -> None:
+        self.ancestors = list(set(self.ancestors or []) | set(state.ancestors or []))
+        self.child_modules = set(self.child_modules) | set(state.child_modules)
+        self.dependencies = list(set(self.dependencies) | set(state.dependencies))
+
+        dep_line_map = {k: -v for k, v in state.dep_line_map.items()}
+        dep_line_map.update(self.dep_line_map)
+        self.dep_line_map = dep_line_map
+
+        priorities = {k: -v for k, v in state.priorities.items()}
+        priorities.update(self.priorities)
+        self.priorities = priorities
+
+        self.dep_line_map.update({k: -v for k, v in state.dep_line_map.items()})
+        if self.tree is not None and state.tree is not None:
+            MergeFiles(self.tree, state.tree, errors, self.xpath, self.id).run()
+
 
 # Module import and diagnostic glue
 
@@ -2199,6 +2219,8 @@ def find_module_and_diagnose(manager: BuildManager,
         file_id = '__builtin__'
     path = find_module_simple(file_id, manager)
     if path:
+        if isinstance(path, tuple):
+            path = path[-1]
         # For non-stubs, look at options.follow_imports:
         # - normal (default) -> fully analyze
         # - silent -> analyze but silence errors
@@ -2274,7 +2296,7 @@ def exist_added_packages(suppressed: List[str],
     return False
 
 
-def find_module_simple(id: str, manager: BuildManager) -> Optional[str]:
+def find_module_simple(id: str, manager: BuildManager) -> Union[Optional[str], Tuple[str, ...]]:
     """Find a filesystem path for module `id` or `None` if not found."""
     t0 = time.time()
     x = manager.find_module_cache.find_module(id)
@@ -2535,8 +2557,16 @@ def load_graph(sources: List[BuildSource], manager: BuildManager,
     # Seed the graph with the initial root sources.
     for bs in sources:
         try:
+            stub_state = None
+            if bs.merge_with:
+                mw = bs.merge_with
+                stub_state = State(id=mw.module, path=mw.path, source=mw.text,
+                                   manager=manager, root_source=True)
+                stub_state.parse_file()
             st = State(id=bs.module, path=bs.path, source=bs.text, manager=manager,
                        root_source=True)
+            if stub_state is not None:
+                st.merge_with(stub_state, manager.errors)
         except ModuleNotFound:
             continue
         if st.id in graph:

mypy/find_sources.py

@@ -1,8 +1,9 @@
 """Routines for finding the sources that mypy will check"""
+from itertools import tee, filterfalse
 
 import os.path
 
-from typing import List, Sequence, Set, Tuple, Optional, Dict
+from typing import List, Sequence, Set, Tuple, Optional, Dict, Callable, Iterable
 
 from mypy.modulefinder import BuildSource, PYTHON_EXTENSIONS
 from mypy.fscache import FileSystemCache
@@ -19,53 +20,76 @@ class InvalidSourceList(Exception):
     """Exception indicating a problem in the list of sources given to mypy."""
 
 
-def create_source_list(files: Sequence[str], options: Options,
-                       fscache: Optional[FileSystemCache] = None,
-                       allow_empty_dir: bool = False) -> List[BuildSource]:
+def partition(
+    pred: Callable[[str], bool], iterable: Iterable[str]
+) -> Tuple[Iterable[str], Iterable[str]]:
+    """Use a predicate to partition entries into false entries and true entries"""
+    t1, t2 = tee(iterable)
+    return filterfalse(pred, t1), filter(pred, t2)
+
+
+def create_source_list(
+    files: Sequence[str],
+    options: Options,
+    fscache: Optional[FileSystemCache] = 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 FileSystemCache()
-    finder = SourceFinder(fscache)
-
+    finder = SourceFinder(fscache, options.merge_stub_into_src)
     targets = []
-    for f in files:
-        if f.endswith(PY_EXTENSIONS):
+    found_targets = set()  # type: Set[str]
+    other, stubs = partition(lambda v: v.endswith(".pyi"), files)
+    source_then_stubs = list(other) + list(stubs)
+    for f in source_then_stubs:
+        if f in found_targets:
+            continue
+        base, ext = os.path.splitext(f)
+        found_targets.add(f)
+        if ext in PY_EXTENSIONS:
             # Can raise InvalidSourceList if a directory doesn't have a valid module name.
             name, base_dir = finder.crawl_up(os.path.normpath(f))
-            targets.append(BuildSource(f, name, None, base_dir))
+            merge_stub = None  # type: Optional[BuildSource]
+            if options.merge_stub_into_src and ext == ".py":
+                stub_file = "{}.pyi".format(base)
+                if os.path.exists(stub_file):
+                    merge_stub = BuildSource(stub_file, name, None, base_dir)
+                    found_targets.add(stub_file)
+            targets.append(BuildSource(f, name, None, base_dir, merge_with=merge_stub))
         elif fscache.isdir(f):
             sub_targets = finder.expand_dir(os.path.normpath(f))
             if not sub_targets and not allow_empty_dir:
-                raise InvalidSourceList("There are no .py[i] files in directory '{}'"
-                                        .format(f))
+                raise InvalidSourceList("There are no .py[i] files in directory '{}'".format(f))
             targets.extend(sub_targets)
         else:
             mod = os.path.basename(f) if options.scripts_are_modules else None
             targets.append(BuildSource(f, mod, None))
     return targets
 
 
-def keyfunc(name: str) -> Tuple[int, str]:
+PY_MAP = {k: i for i, k in enumerate(PY_EXTENSIONS)}
+
+
+def keyfunc(name: str) -> Tuple[str, int]:
     """Determines sort order for directory listing.
 
     The desirable property is foo < foo.pyi < foo.py.
     """
     base, suffix = os.path.splitext(name)
-    for i, ext in enumerate(PY_EXTENSIONS):
-        if suffix == ext:
-            return (i, base)
-    return (-1, name)
+    return base, PY_MAP.get(suffix, -1)
 
 
 class SourceFinder:
-    def __init__(self, fscache: FileSystemCache) -> None:
+    def __init__(self, fscache: FileSystemCache, merge_stub_into_src: bool) -> None:
         self.fscache = fscache
         # A cache for package names, mapping from directory path to module id and base dir
         self.package_cache = {}  # type: Dict[str, Tuple[str, str]]
+        self.merge_stub_into_src = merge_stub_into_src
 
-    def expand_dir(self, arg: str, mod_prefix: str = '') -> List[BuildSource]:
+    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:
@@ -79,27 +103,46 @@ def expand_dir(self, arg: str, mod_prefix: str = '') -> List[BuildSource]:
             sources.append(BuildSource(f, mod_prefix.rstrip('.'), None, base_dir))
         names = self.fscache.listdir(arg)
         names.sort(key=keyfunc)
-        for name in names:
-            # Skip certain names altogether
-            if (name == '__pycache__' or name == 'py.typed'
-                    or name.startswith('.')
-                    or name.endswith(('~', '.pyc', '.pyo'))):
-                continue
-            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__':
+        name_iter = iter(names)
+        try:
+            name = next(name_iter, None)
+            while name is not None:
+                # Skip certain names altogether
+                if (name == '__pycache__' or name == 'py.typed'
+                        or name.startswith('.')
+                        or name.endswith(('~', '.pyc', '.pyo'))):
                     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, base_dir)
-                    sources.append(src)
-        return sources
+                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)
+                    name = next(name_iter)
+                else:
+                    base, suffix = os.path.splitext(name)
+                    name = next(name_iter, None)
+                    if base == '__init__':
+                        continue
+                    if base not in seen and '.' not in base and suffix in PY_EXTENSIONS:
+                        seen.add(base)
+                        if name is None:
+                            next_base, next_suffix = None, None
+                        else:
+                            next_base, next_suffix = os.path.splitext(name)
+                        src = BuildSource(path, mod_prefix + base, None, base_dir)
+                        if self.merge_stub_into_src is True and next_base is not None \
+                                and next_base == base and name is not None:
+                            merge_with = src
+                            src = BuildSource(path=os.path.join(arg, name),
+                                              module=mod_prefix + next_base,
+                                              merge_with=merge_with,
+                                              text=None,
+                                              base_dir=base_dir)
+                        sources.append(src)
+            return sources
+        except StopIteration:
+            return sources
 
     def crawl_up(self, arg: str) -> Tuple[str, str]:
         """Given a .py[i] filename, return module and base directory

mypy/main.py

@@ -631,6 +631,10 @@ def add_invertible_flag(flag: str,
         '--find-occurrences', metavar='CLASS.MEMBER',
         dest='special-opts:find_occurrences',
         help="Print out all usages of a class member (experimental)")
+    add_invertible_flag('--merge-stub-into-src', default=False, strict_flag=False,
+                        help="when a stub and source file is in the same folder with same name "
+                             "merge the stub file into the source file, and lint the source file",
+                        group=other_group)
 
     if server_options:
         # TODO: This flag is superfluous; remove after a short transition (2018-03-16)

mypy/modulefinder.py

@@ -10,7 +10,7 @@
 import subprocess
 import sys
 
-from typing import Dict, List, NamedTuple, Optional, Set, Tuple
+from typing import Any, Dict, List, NamedTuple, Optional, Set, Tuple, Union
 
 MYPY = False
 if MYPY:
@@ -41,17 +41,32 @@ 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], base_dir: Optional[str] = None,
+                 merge_with: Optional['BuildSource'] = 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
         self.base_dir = base_dir  # Directory where the package is rooted (e.g. 'xxx/yyy')
+        self.merge_with = merge_with
 
     def __repr__(self) -> str:
         return '<BuildSource path=%r module=%r has_text=%s>' % (self.path,
                                                                 self.module,
                                                                 self.text is not None)
 
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, BuildSource):
+            return False
+        return (self.path, self.module, self.text, self.base_dir, self.merge_with) == (
+            other.path,
+            other.module,
+            other.text,
+            other.base_dir,
+            other.merge_with,
+        )
+
+    def __ne__(self, other: Any) -> bool:
+        return not self.__eq__(other)
 
 class FindModuleCache:
     """Module finder with integrated cache.
@@ -232,13 +247,18 @@ def _find_module(self, id: str) -> Optional[str]:
                 elif self.options and self.options.namespace_packages and fscache.isdir(base_path):
                     near_misses.append(base_path)
             # No package, look for module.
+            paths = []
             for extension in PYTHON_EXTENSIONS:
                 path = base_path + extension
                 if fscache.isfile_case(path):
                     if verify and not verify_module(fscache, id, path):
                         near_misses.append(path)
                         continue
-                    return path
+                    paths.append(path)
+            if len(paths) == 1:
+                return paths[0]
+            elif len(paths) > 1:
+                return tuple(paths)
 
         # In namespace mode, re-check those entries that had 'verify'.
         # Assume search path entries xxx, yyy and zzz, and we're
@@ -276,7 +296,10 @@ def find_modules_recursive(self, module: str) -> List[BuildSource]:
         module_path = self.find_module(module)
         if not module_path:
             return []
-        result = [BuildSource(module_path, module, None)]
+        merge_with = None
+        if isinstance(module_path, tuple):
+            module_path = module_path
+        result = [BuildSource(module_path, module, None, merge_with=merge_with)]
         if module_path.endswith(('__init__.py', '__init__.pyi')):
             # Subtle: this code prefers the .pyi over the .py if both
             # exists, and also prefers packages over modules if both x/

mypy/options.py

@@ -241,6 +241,8 @@ def __init__(self) -> None:
         # Don't properly free objects on exit, just kill the current process.
         self.fast_exit = False
 
+        self.merge_stub_into_src = False
+
     def snapshot(self) -> object:
         """Produce a comparable snapshot of this Option"""
         # Under mypyc, we don't have a __dict__, so we need to do worse things.