Fix relative imports and improve private name handling

JukkaL · JukkaL · commit 6b3334d32adc · 2019-11-07T12:05:27.000Z
diff --git a/mypy/stubgen.py b/mypy/stubgen.py
@@ -484,6 +484,7 @@ def __init__(self,
 
     def visit_mypy_file(self, o: MypyFile) -> None:
         self.module = o.fullname()  # Current module being processed
+        self.path = o.path
         self.defined_names = find_defined_names(o)
         self.referenced_names = find_referenced_names(o)
         typing_imports = ["Any", "Optional", "TypeVar"]
@@ -833,6 +834,14 @@ def visit_import_from(self, o: ImportFrom) -> None:
         exported_names = set()  # type: Set[str]
         import_names = []
         module, relative = self.translate_module_name(o.id, o.relative)
+        if self.module:
+            full_module, ok = mypy.util.correct_relative_import(
+                self.module, relative, module, self.path.endswith('.__init__.py')
+            )
+            if not ok:
+                full_module = module
+        else:
+            full_module = module
         if module == '__future__':
             return  # Not preserved
         for name, as_name in o.names:
@@ -842,18 +851,24 @@ def visit_import_from(self, o: ImportFrom) -> None:
                 continue
             exported = False
             if as_name is None and self.module and (self.module + '.' + name) in EXTRA_EXPORTED:
+                # Special case certain names that should be exported, against our general rules.
                 exported = True
-            if (as_name is None and name not in self.referenced_names and not self._all_
+            is_private = self.is_private_name(name, full_module + '.' + name)
+            if (as_name is None
+                    and name not in self.referenced_names
+                    and not self._all_
+                    and not is_private
                     and module not in ('abc', 'typing', 'asyncio')):
                 # An imported name that is never referenced in the module is assumed to be
                 # 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 = module.split('.')[0]
+            top_level = full_module.split('.')[0]
             if (as_name is None
                     and not self.export_less
                     and not self._all_
                     and self.module
+                    and not is_private
                     and top_level in (self.module.split('.')[0],
                                       '_' + self.module.split('.')[0])):
                 # Export imports from the same package, since we can't reliably tell whether they
diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test
@@ -1967,3 +1967,51 @@ c: C
 # p/a.pyi
 class A: ...
 class B: ...
+
+[case testExportInternalImportsByDefaultUsingRelativeImport]
+# modules: p.a
+
+[file p/__init__.py]
+
+[file p/a.py]
+from .b import f
+f()
+
+[file p/b.py]
+def f(): pass
+
+[out]
+from .b import f as f
+
+[case testExportInternalImportsByDefaultSkipPrivate]
+# modules: p.a
+
+[file p/__init__.py]
+
+[file p/a.py]
+from .b import _f, _g as _g
+_f()
+
+[file p/b.py]
+def _f(): pass
+def _g(): pass
+x = 0
+
+[out]
+from .b import _f
+
+[case testExportInternalImportsByDefaultIncludePrivate]
+# flags: --include-private
+# modules: p.a
+
+[file p/__init__.py]
+
+[file p/a.py]
+from .b import _f
+_f()
+
+[file p/b.py]
+def _f(): pass
+
+[out]
+from .b import _f as _f
