Repository: python/mypy

Issue:
stubgen: wrong import statements for base module when importing submodule
There's a bug with the import statements generated in .pyi files.

Consider the folder structure:
```
somefolder/
    test.py
    foo_mod/
        __init__.py
        bar_mod.py
```
with file contents:
```python
# test.py
import foo_mod
import foo_mod.bar_mod as bar_mod

class baz(foo_mod.foo, bar_mod.bar):
    pass
```
```python
# foo_mod/__init__.py
class foo:
    pass
```
```python
# foo_mod/bar_mod.py
class bar:
    pass
```
Then running `stubgen test.py` whilst in `somefolder` generates an `out/test.pyi` file with contents:
```python
# out/test.pyi
# Stubs for test (Python 3)
#
# NOTE: This dynamically typed stub was automatically generated by stubgen.

import foo_mod.bar_mod as bar_mod
import foo_mod.bar_mod

class baz(foo_mod.foo, bar_mod.bar): ...
```
Note how it says `foo_mod.bar_mod`, rather than just `foo_mod` as expected.

This issue generalises to subsubpackages as well (and presumably further). That is,
```python
# test.py
import foo_mod
import foo_mod.bar_mod as bar_mod
import foo_mod.bar_mod.quux_mod as quux_mod
...
```
generates
```python
# out/test.pyi
import foo_mod.bar_mod.quux_mod as bar_mod
import foo_mod.bar_mod.quux_mod
import foo_mod.bar_mod.quux_mod as quux_mod
...
```

Tested on Windows and Unix, with both mypy version 0.701 and the latest on the GitHub master branch, with Python 3.6.6. (At least, that's the version of Python which is bound to `python` on the command line, so presumably that's what stubgen is using?)

Possibly related? https://github.com/python/mypy/issues/6831


Discussion:
It seems that it's whichever import statement is last that is used. e.g.
```python
# test.py
import foo_mod.bar_mod as bar_mod
import foo_mod
...
```
generates
```python
# out/test.pyi
import foo_mod as bar_mod
import foo_mod
...
```
I think this is a separate issue from #6831.

Patch:
diff --git a/mypy/stubgen.py b/mypy/stubgen.py
--- a/mypy/stubgen.py
+++ b/mypy/stubgen.py
@@ -496,7 +496,9 @@ def add_import(self, module: str, alias: str | None = None) -> None:
                 name = name.rpartition(".")[0]
 
     def require_name(self, name: str) -> None:
-        self.required_names.add(name.split(".")[0])
+        while name not in self.direct_imports and "." in name:
+            name = name.rsplit(".", 1)[0]
+        self.required_names.add(name)
 
     def reexport(self, name: str) -> None:
         """Mark a given non qualified name as needed in __all__.
@@ -516,7 +518,10 @@ def import_lines(self) -> list[str]:
         # be imported from it. the names can also be alias in the form 'original as alias'
         module_map: Mapping[str, list[str]] = defaultdict(list)
 
-        for name in sorted(self.required_names):
+        for name in sorted(
+            self.required_names,
+            key=lambda n: (self.reverse_alias[n], n) if n in self.reverse_alias else (n, ""),
+        ):
             # If we haven't seen this name in an import statement, ignore it
             if name not in self.module_for:
                 continue
@@ -540,7 +545,7 @@ def import_lines(self) -> list[str]:
                     assert "." not in name  # Because reexports only has nonqualified names
                     result.append(f"import {name} as {name}\n")
                 else:
-                    result.append(f"import {self.direct_imports[name]}\n")
+                    result.append(f"import {name}\n")
 
         # Now generate all the from ... import ... lines collected in module_map
         for module, names in sorted(module_map.items()):
@@ -595,7 +600,7 @@ def visit_name_expr(self, e: NameExpr) -> None:
         self.refs.add(e.name)
 
     def visit_instance(self, t: Instance) -> None:
-        self.add_ref(t.type.fullname)
+        self.add_ref(t.type.name)
         super().visit_instance(t)
 
     def visit_unbound_type(self, t: UnboundType) -> None:
@@ -614,7 +619,10 @@ def visit_callable_type(self, t: CallableType) -> None:
         t.ret_type.accept(self)
 
     def add_ref(self, fullname: str) -> None:
-        self.refs.add(fullname.split(".")[-1])
+        self.refs.add(fullname)
+        while "." in fullname:
+            fullname = fullname.rsplit(".", 1)[0]
+            self.refs.add(fullname)
 
 
 class StubGenerator(mypy.traverser.TraverserVisitor):
@@ -1295,6 +1303,7 @@ def visit_import_from(self, o: ImportFrom) -> None:
             if (
                 as_name is None
                 and name not in self.referenced_names
+                and not any(n.startswith(name + ".") for n in self.referenced_names)
                 and (not self._all_ or name in IGNORED_DUNDERS)
                 and not is_private
                 and module not in ("abc", "asyncio") + TYPING_MODULE_NAMES
@@ -1303,14 +1312,15 @@ def visit_import_from(self, o: ImportFrom) -> None:
                 # exported, unless there is an explicit __all__. Note that we need to special
                 # case 'abc' since some references are deleted during semantic analysis.
                 exported = True
-            top_level = full_module.split(".")[0]
+            top_level = full_module.split(".", 1)[0]
+            self_top_level = self.module.split(".", 1)[0]
             if (
                 as_name is None
                 and not self.export_less
                 and (not self._all_ or name in IGNORED_DUNDERS)
                 and self.module
                 and not is_private
-                and top_level in (self.module.split(".")[0], "_" + self.module.split(".")[0])
+                and top_level in (self_top_level, "_" + self_top_level)
             ):
                 # Export imports from the same package, since we can't reliably tell whether they
                 # are part of the public API.


Test patch:
diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test
--- a/test-data/unit/stubgen.test
+++ b/test-data/unit/stubgen.test
@@ -2772,9 +2772,9 @@ y: b.Y
 z: p.a.X
 
 [out]
+import p.a
 import p.a as a
 import p.b as b
-import p.a
 
 x: a.X
 y: b.Y
@@ -2787,7 +2787,7 @@ from p import a
 x: a.X
 
 [out]
-from p import a as a
+from p import a
 
 x: a.X
 
@@ -2809,7 +2809,7 @@ from p import a
 x: a.X
 
 [out]
-from p import a as a
+from p import a
 
 x: a.X
 
@@ -2859,6 +2859,60 @@ import p.a
 x: a.X
 y: p.a.Y
 
+[case testNestedImports]
+import p
+import p.m1
+import p.m2
+
+x: p.X
+y: p.m1.Y
+z: p.m2.Z
+
+[out]
+import p
+import p.m1
+import p.m2
+
+x: p.X
+y: p.m1.Y
+z: p.m2.Z
+
+[case testNestedImportsAliased]
+import p as t
+import p.m1 as pm1
+import p.m2 as pm2
+
+x: t.X
+y: pm1.Y
+z: pm2.Z
+
+[out]
+import p as t
+import p.m1 as pm1
+import p.m2 as pm2
+
+x: t.X
+y: pm1.Y
+z: pm2.Z
+
+[case testNestedFromImports]
+from p import m1
+from p.m1 import sm1
+from p.m2 import sm2
+
+x: m1.X
+y: sm1.Y
+z: sm2.Z
+
+[out]
+from p import m1
+from p.m1 import sm1
+from p.m2 import sm2
+
+x: m1.X
+y: sm1.Y
+z: sm2.Z
+
 [case testOverload_fromTypingImport]
 from typing import Tuple, Union, overload