Support importing types from nested packages

mypy/build.py

@@ -841,7 +841,8 @@ def stats_summary(self) -> Mapping[str, object]:
 
 
 # Package dirs are a two-tuple of path to search and whether to verify the module
-PackageDirs = List[Tuple[str, bool]]
+OnePackageDir = Tuple[str, bool]
+PackageDirs = List[OnePackageDir]
 
 
 class FindModuleCache:
@@ -896,6 +897,15 @@ def find_module(self, id: str, search_paths: SearchPaths,
             self.results[key] = self._find_module(id, search_paths, python_executable)
         return self.results[key]
 
+    def _find_module_non_stub_helper(self, components: List[str],
+                                     pkg_dir: str) -> Optional[OnePackageDir]:
+        dir_path = pkg_dir
+        for index, component in enumerate(components):
+            dir_path = os.path.join(dir_path, component)
+            if self.fscache.isfile(os.path.join(dir_path, 'py.typed')):
+                return os.path.join(pkg_dir, *components[:-1]), index == 0
+        return None
+
     def _find_module(self, id: str, search_paths: SearchPaths,
                      python_executable: Optional[str]) -> Optional[str]:
         fscache = self.fscache
@@ -910,12 +920,11 @@ def _find_module(self, id: str, search_paths: SearchPaths,
 
         # We have two sets of folders so that we collect *all* stubs folders and
         # put them in the front of the search path
-        third_party_inline_dirs = []
-        third_party_stubs_dirs = []
+        third_party_inline_dirs = []  # type: PackageDirs
+        third_party_stubs_dirs = []  # type: PackageDirs
         # Third-party stub/typed packages
         for pkg_dir in search_paths.package_path:
             stub_name = components[0] + '-stubs'
-            typed_file = os.path.join(pkg_dir, components[0], 'py.typed')
             stub_dir = os.path.join(pkg_dir, stub_name)
             if fscache.isdir(stub_dir):
                 stub_typed_file = os.path.join(stub_dir, 'py.typed')
@@ -935,9 +944,9 @@ def _find_module(self, id: str, search_paths: SearchPaths,
                         third_party_stubs_dirs.append((path, False))
                     else:
                         third_party_stubs_dirs.append((path, True))
-            elif fscache.isfile(typed_file):
-                path = os.path.join(pkg_dir, dir_chain)
-                third_party_inline_dirs.append((path, True))
+            non_stub_match = self._find_module_non_stub_helper(components, pkg_dir)
+            if non_stub_match:
+                third_party_inline_dirs.append(non_stub_match)
         if self.options and self.options.use_builtins_fixtures:
             # Everything should be in fixtures.
             third_party_inline_dirs.clear()
@@ -2357,11 +2366,13 @@ def find_module_and_diagnose(manager: BuildManager,
             if not (options.ignore_missing_imports or in_partial_package(id, manager)):
                 module_not_found(manager, caller_line, caller_state, id)
             raise ModuleNotFound
-        else:
+        elif root_source:
             # If we can't find a root source it's always fatal.
             # TODO: This might hide non-fatal errors from
             # root sources processed earlier.
             raise CompileError(["mypy: can't find module '%s'" % id])
+        else:
+            raise ModuleNotFound
 
 
 def exist_added_packages(suppressed: List[str],

mypy/test/testpep561.py

@@ -21,6 +21,17 @@
 reveal_type(a)
 """
 
+NAMESPACE_PROGRAM = """
+from typedpkg_nested.nested_package.nested_module import nested_func
+from typedpkg_namespace.alpha.alpha_module import alpha_func
+
+nested_func("abc")
+alpha_func(False)
+
+nested_func(False)
+alpha_func(2)
+"""
+
 
 def check_mypy_run(cmd_line: List[str],
                    python_executable: str = sys.executable,
@@ -95,13 +106,24 @@ def setUp(self) -> None:
         self.tempfile = os.path.join(self.temp_file_dir.name, 'simple.py')
         with open(self.tempfile, 'w+') as file:
             file.write(SIMPLE_PROGRAM)
+        self.namespace_tempfile = os.path.join(self.temp_file_dir.name, 'namespace_program.py')
+        with open(self.namespace_tempfile, 'w+') as file:
+            file.write(NAMESPACE_PROGRAM)
+
         self.msg_dne = \
             "{}:3: error: Module 'typedpkg' has no attribute 'dne'\n".format(self.tempfile)
         self.msg_list = \
             "{}:5: error: Revealed type is 'builtins.list[builtins.str]'\n".format(self.tempfile)
         self.msg_tuple = \
             "{}:5: error: Revealed type is 'builtins.tuple[builtins.str]'\n".format(self.tempfile)
 
+        self.namespace_msg_bool_str = (
+            '{0}:8: error: Argument 1 to "nested_func" has incompatible type "bool"; '
+            'expected "str"\n'.format(self.namespace_tempfile))
+        self.namespace_msg_int_bool = (
+            '{0}:9: error: Argument 1 to "alpha_func" has incompatible type "int"; '
+            'expected "bool"\n'.format(self.namespace_tempfile))
+
     def tearDown(self) -> None:
         self.temp_file_dir.cleanup()
 
@@ -192,6 +214,18 @@ def test_typedpkg_editable(self) -> None:
                 venv_dir=venv_dir,
             )
 
+    def test_nested_and_namespace(self) -> None:
+        with self.virtualenv() as venv:
+            venv_dir, python_executable = venv
+            self.install_package('typedpkg_nested', python_executable)
+            self.install_package('typedpkg_namespace-alpha', python_executable)
+            check_mypy_run(
+                [self.namespace_tempfile],
+                python_executable,
+                expected_out=self.namespace_msg_bool_str + self.namespace_msg_int_bool,
+                venv_dir=venv_dir,
+            )
+
 
 if __name__ == '__main__':
     main()

test-data/packages/typedpkg_namespace-alpha/setup.py

@@ -0,0 +1,10 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='typedpkg_namespace.alpha',
+    version='1.0.0',
+    packages=find_packages(),
+    namespace_packages=['typedpkg_namespace'],
+    zip_safe=False,
+    package_data={'typedpkg_namespace.alpha': ['py.typed']}
+)

test-data/packages/typedpkg_namespace-alpha/typedpkg_namespace/__init__.py

@@ -0,0 +1,2 @@
+# namespace pkg
+__import__("pkg_resources").declare_namespace(__name__)

test-data/packages/typedpkg_namespace-alpha/typedpkg_namespace/alpha/alpha_module.py

@@ -0,0 +1,2 @@
+def alpha_func(a: bool) -> bool:
+    return not a

test-data/packages/typedpkg_nested/setup.py

@@ -0,0 +1,9 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='typedpkg_nested',
+    version='1.0.0',
+    packages=find_packages(),
+    zip_safe=False,
+    package_data={'typedpkg_nested.nested_package': ['py.typed']}
+)

test-data/packages/typedpkg_nested/typedpkg_nested/nested_package/nested_module.py

@@ -0,0 +1,2 @@
+def nested_func(a: str) -> str:
+    return a + " nested"