Fix a bug in pytype's import resolution for relative imports.

We made the (erroneous) assumption that "from . import X" can only import a
module, which was causing trouble for typeshed. See
python/typeshed#5762 (comment) for
context. Fixing this exposed another bug where load_pytd.Module.is_package
didn't work correctly for stubs that ship with pytype, which I also fixed.

PiperOrigin-RevId: 388369899

Author: rchen152

pytype/load_pytd.py

@@ -384,11 +384,18 @@ def _parse_predefined(self, pytd_subdir, module, as_package=False):
     return ast
 
   def get_builtin(self, subdir, module_name):
+    """Load a stub that ships with pytype."""
     builtin_dir = file_utils.get_versioned_path(subdir, self.python_version)
     mod = self._parse_predefined(builtin_dir, module_name)
-    if not mod:
+    # For stubs in pytype's stubs/ directory, we use the module name prefixed
+    # with "pytd:" for the filename. Package filenames need an "/__init__.pyi"
+    # suffix for Module.is_package to recognize them.
+    if mod:
+      filename = module_name
+    else:
       mod = self._parse_predefined(builtin_dir, module_name, as_package=True)
-    return mod
+      filename = os.path.join(module_name, "__init__.pyi")
+    return filename, mod
 
 
 class Loader:
@@ -695,9 +702,9 @@ def _load_builtin(self, subdir, module_name, third_party_only=False):
     """Load a pytd/pyi that ships with pytype or typeshed."""
     # Try our own type definitions first.
     if not third_party_only:
-      mod_ast = self._builtin_loader.get_builtin(subdir, module_name)
+      filename, mod_ast = self._builtin_loader.get_builtin(subdir, module_name)
       if mod_ast:
-        return self.load_file(filename=self.PREFIX + module_name,
+        return self.load_file(filename=self.PREFIX + filename,
                               module_name=module_name, mod_ast=mod_ast)
     if self.use_typeshed:
       return self._load_typeshed_builtin(subdir, module_name)

pytype/load_pytd_test.py

@@ -483,6 +483,16 @@ def f(x: T) -> T: ...
       from b import *
     """)
 
+  def test_import_class_from_parent_module(self):
+    with file_utils.Tempdir() as d:
+      d.create_file("foo/__init__.pyi", "class Foo: ...")
+      d.create_file("foo/bar.pyi", """
+        from . import Foo
+        class Bar(Foo): ...
+      """)
+      loader = load_pytd.Loader(None, self.python_version, pythonpath=[d.path])
+      loader.import_name("foo.bar")
+
 
 class ImportTypeMacroTest(_LoaderTest):
 

pytype/pyi/modules.py

@@ -1,8 +1,7 @@
 """Handling of module and package related details."""
 
-from typing import Any
-
 import dataclasses
+from typing import Any
 
 from pytype import file_utils
 from pytype import module_utils
@@ -91,20 +90,14 @@ def process_from_import(self, from_package, item):
     else:
       name = new_name = item
     qualified_name = self.qualify_name("%s.%s" % (from_package, name))
-    if (from_package in ["__PACKAGE__", "__PARENT__"]
-        and isinstance(item, str)):
-      # This will always be a simple module import (from . cannot import a
-      # NamedType, and without 'as' the name will not be reexported).
-      t = pytd.Module(name=new_name, module_name=qualified_name)
-    else:
-      # We should ideally not need this check, but we have typing
-      # special-cased in some places.
-      if not qualified_name.startswith("typing.") and name != "*":
-        # Mark this as an externally imported type, so that AddNamePrefix
-        # does not prefix it with the current package name.
-        qualified_name = (parser_constants.EXTERNAL_NAME_PREFIX +
-                          qualified_name)
-      t = pytd.NamedType(qualified_name)
+    # We should ideally not need this check, but we have typing
+    # special-cased in some places.
+    if not qualified_name.startswith("typing.") and name != "*":
+      # Mark this as an externally imported type, so that AddNamePrefix
+      # does not prefix it with the current package name.
+      qualified_name = (parser_constants.EXTERNAL_NAME_PREFIX +
+                        qualified_name)
+    t = pytd.NamedType(qualified_name)
     if name == "*":
       # A star import is stored as
       # 'imported_mod.* = imported_mod.*'. The imported module needs to be