Some more cases

Author: m-novikov

mypy/build.py

@@ -860,35 +860,41 @@ def find_module(self, id: str) -> Optional[str]:
             return None
 
     def find_modules_recursive(self, module: str) -> List[BuildSource]:
+        hits = set()  # type: Set[str]
+        result = []  # type: List[BuildSource]
+        for src in self._find_modules_recursive(module):
+            if src.module not in hits:
+                hits.add(src.module)
+                result.append(src)
+        return result
+
+    def _find_modules_recursive(self, module: str) -> List[BuildSource]:
         module_paths = self._find_module(module)
 
-        result = []
-        hits = set()  # type: Set[str]
+        srcs = []  # type: List[BuildSource]
         for path in module_paths:
             if is_module_path(path) or is_pkg_path(path):
-                if module not in hits:
-                    result.append(BuildSource(path, module, None))
-                    hits.add(module)
+                srcs.append(BuildSource(path, module, None))
 
             if is_pkg_path(path):
                 path = dirname(path)
-                result += self._traverse_package(module, path)
+                for submodule in self._find_submodules(module, path):
+                    srcs += self._find_modules_recursive(submodule)
             elif is_namespace_path(path):
-                result += self._traverse_package(module, path)
+                for submodule in self._find_submodules(module, path):
+                    srcs += self._find_modules_recursive(submodule)
 
-        return result
+        return srcs
 
-    def _traverse_package(self, module, path) -> List[BuildSource]:
-        result = []  # type: List[BuildSource]
+    def _find_submodules(self, module, path) -> Iterator[str]:
         for item in list_dir(path):
             if item == '__init__.py' or item == '__init__.pyi':
                 continue
 
             if item.endswith(tuple(PYTHON_EXTENSIONS)):
                 item = item.split('.')[0]
 
-            result += self.find_modules_recursive(module + '.' + item)
-        return result
+            yield module + '.' + item
 
     def _collect_paths(self, paths: List[str], last_comp: str) -> List[str]:
         """

mypy/test/testmodulediscovery.py

@@ -125,3 +125,14 @@ def test_find_modules_recursive_with_namespace(self) -> None:
         m = ModuleDiscovery(['dir1', 'dir2'], namespaces_allowed=True)
         srcs = m.find_modules_recursive('mod')
         assert_equal([s.module for s in srcs], ['mod.a', 'mod.b'])
+
+    def test_find_modules_recursive_with_stubs(self) -> None:
+        self.files = {
+            os.path.join('dir1', 'mod', '__init__.py'),
+            os.path.join('dir1', 'mod', 'a.py'),
+            os.path.join('dir2', 'mod', '__init__.pyi'),
+            os.path.join('dir2', 'mod', 'a.pyi'),
+        }
+        m = ModuleDiscovery(['dir1', 'dir2'], namespaces_allowed=True)
+        srcs = m.find_modules_recursive('mod')
+        assert_equal([s.module for s in srcs], ['mod', 'mod.a'])