add tests for source tree and fix annotations

.gitignore

@@ -19,6 +19,7 @@ dmypy.json
 # Packages
 *.egg
 *.egg-info
+.eggs
 
 # IDEs
 .idea

mypy/build.py

@@ -111,6 +111,12 @@ def __repr__(self) -> str:
                                                                 self.module,
                                                                 self.text is not None)
 
+    def __eq__(self, other: Any) -> bool:
+        return type(self) == type(other) and self.__dict__ == other.__dict__
+
+    def __ne__(self, other):
+        return not self == other
+
 
 class BuildSourceSet:
     """Efficiently test a file's membership in the set of build sources."""
@@ -2562,11 +2568,10 @@ def load_graph(sources: List[BuildSource], manager: BuildManager,
                        root_source=True)
             if bs.merge_with:
                 mw = bs.merge_with
-                src_st = State(id=mw.module, path=mw.path, source=mw.text,
-                               manager=manager, root_source=True)
-                src_st.parse_file()
-                src_st.merge_with(st, manager.errors)
-                st = src_st
+                stub_state = State(id=mw.module, path=mw.path, source=mw.text,
+                                   manager=manager, root_source=True)
+                stub_state.parse_file()
+                st.merge_with(stub_state, manager.errors)
         except ModuleNotFound:
             continue
         if st.id in graph:

mypy/find_sources.py

@@ -24,7 +24,7 @@ def create_source_list(files: Sequence[str], options: Options,
     Raises InvalidSourceList on errors.
     """
     fscache = fscache or FileSystemCache()
-    finder = SourceFinder(fscache)
+    finder = SourceFinder(fscache, options.merge_stub_into_src)
 
     targets = []
     for f in files:
@@ -57,10 +57,11 @@ def keyfunc(name: str) -> Tuple[str, int]:
 
 
 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]:
         """Convert a directory name to a list of sources to build."""
@@ -98,14 +99,15 @@ def expand_dir(self, arg: str, mod_prefix: str = '') -> List[BuildSource]:
                             next_base, next_suffix = None, None
                         else:
                             next_base, next_suffix = os.path.splitext(name)
-                        merge_with = None
-                        if next_base is not None and next_base == base and name is not None:
-                            merge_with = BuildSource(os.path.join(arg, name),
-                                                     mod_prefix + next_base,
-                                                     None,
-                                                     base_dir)
-                        src = BuildSource(path, mod_prefix + base, None, base_dir,
-                                          merge_with=merge_with)
+                        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:

mypy/main.py

@@ -585,6 +585,10 @@ def add_invertible_flag(flag: str,
     other_group.add_argument(
         '--scripts-are-modules', action='store_true',
         help="script x becomes module x instead of __main__")
+    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/options.py

@@ -200,6 +200,8 @@ def __init__(self) -> None:
         self.package_root = []  # type: List[str]
         self.cache_map = {}  # type: Dict[str, Tuple[str, str]]
 
+        self.merge_stub_into_src = False
+
     def snapshot(self) -> object:
         """Produce a comparable snapshot of this Option"""
         d = dict(self.__dict__)

mypy/stub_src_merge.py

@@ -10,7 +10,7 @@
     Tuple,
     Type,
     Union,
-    cast
+    cast,
 )
 
 from mypy.errors import Errors
@@ -32,7 +32,7 @@
     SymbolTable,
     TempNode,
     TypeInfo,
-    Var
+    Var,
 )
 
 
@@ -139,7 +139,7 @@ def merge_assignment(self, src: AssignmentStmt, stub: AssignmentStmt) -> None:
         # already made sure they match
         src.type = stub.type
         src.unanalyzed_type = stub.unanalyzed_type
-        stub_name = self.simple_assignment_name(stub)
+        stub_name = cast(str, self.simple_assignment_name(stub))
         self.check_no_explicit_default(stub.rvalue, stub, stub_name)
 
     def merge_func_definition(self, src: FuncDef, stub: FuncDef) -> None:
@@ -295,7 +295,7 @@ def decorator_argument_checks(self, src: CallExpr, stub: CallExpr) -> None:
                 )
                 break
         for name, default_node in zip(stub.arg_names, stub.args):
-            self.check_no_explicit_default(default_node, src, name)
+            self.check_no_explicit_default(default_node, src, cast(str, name))
 
     def check_no_explicit_default(
         self, default_node: Node, node: Node, name: str

mypy/test/test_source_finder.py

@@ -0,0 +1,108 @@
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Tuple
+
+import pytest
+from py._path.local import LocalPath
+
+from mypy.build import BuildSource
+from mypy.find_sources import create_source_list
+from mypy.options import Options
+
+MergeFinder = Callable[[Dict[str, Any], bool], Tuple[Path, List[BuildSource]]]
+
+
+def create_files(folder: Path, content: Dict[str, Any]):
+    folder.mkdir(exist_ok=True)
+    for key, value in content.items():
+        if key.endswith(".py") or key.endswith(".pyi"):
+            (folder / key).write_text(value)
+        elif isinstance(value, dict):
+            create_files(folder / key, value)
+
+
+@pytest.fixture()
+def merge_finder(tmpdir: LocalPath) -> MergeFinder:
+    def _checker(
+        content: Dict[str, Any], merge: bool
+    ) -> Tuple[Path, List[BuildSource]]:
+        options = Options()
+        options.merge_stub_into_src = merge
+        test_dir = str(tmpdir)
+        create_files(Path(test_dir), content)
+        targets = create_source_list(files=[test_dir], options=options)
+        return Path(str(tmpdir)), targets
+
+    return _checker
+
+
+def test_source_finder_merge(merge_finder: MergeFinder) -> None:
+    base_dir, found = merge_finder({"a.py": "", "a.pyi": ""}, True)
+    assert found == [
+        BuildSource(
+            path=str(base_dir / "a.py"),
+            base_dir=str(base_dir),
+            module="a",
+            text=None,
+            merge_with=BuildSource(
+                path=str(base_dir / "a.pyi"),
+                base_dir=str(base_dir),
+                module="a",
+                text=None,
+            ),
+        )
+    ]
+
+
+def test_source_finder_merge_sub_folder(merge_finder: MergeFinder) -> None:
+    base_dir, found = merge_finder(
+        {"pkg": {"a.py": "", "a.pyi": "", "__init__.py": ""}}, True
+    )
+    assert found == [
+        BuildSource(
+            path=str(base_dir / "pkg" / "__init__.py"),
+            base_dir=str(base_dir),
+            module="pkg",
+            text=None,
+        ),
+        BuildSource(
+            path=str(base_dir / "pkg" / "a.py"),
+            base_dir=str(base_dir),
+            module="pkg.a",
+            text=None,
+            merge_with=BuildSource(
+                path=str(base_dir / "pkg" / "a.pyi"),
+                base_dir=str(base_dir),
+                module="pkg.a",
+                text=None,
+            ),
+        ),
+    ]
+
+
+def test_source_finder_no_merge(merge_finder: MergeFinder) -> None:
+    base_dir, found = merge_finder({"a.py": "", "a.pyi": ""}, False)
+    assert found == [
+        BuildSource(
+            path=str(base_dir / "a.pyi"), base_dir=str(base_dir), module="a", text=None
+        )
+    ]
+
+
+@pytest.mark.parametrize("merge", [True, False])
+def test_source_finder_merge_just_source(merge_finder, merge):
+    base_dir, found = merge_finder({"a.py": ""}, merge=merge)
+    assert found == [
+        BuildSource(
+            path=str(base_dir / "a.py"), base_dir=str(base_dir), module="a", text=None
+        )
+    ]
+
+
+@pytest.mark.parametrize("merge", [True, False])
+def test_source_finder_merge_just_stub(merge_finder, merge):
+    base_dir, found = merge_finder({"a.pyi": ""}, merge=merge)
+    assert found == [
+        BuildSource(
+            path=str(base_dir / "a.pyi"), base_dir=str(base_dir), module="a", text=None
+        )
+    ]