rewrite: add support for "import module.submodule".

rchen152 · rchen152 · commit 7d4ce24a2954 · 2024-04-11T15:10:08.000-07:00
PiperOrigin-RevId: 623629676
diff --git a/pytype/rewrite/abstract/classes_test.py b/pytype/rewrite/abstract/classes_test.py
@@ -58,5 +58,23 @@ def test_get_attribute(self):
     self.assertEqual(instance.get_attribute('x'), self.ctx.consts[3])
 
 
+class ModuleTest(test_utils.ContextfulTestBase):
+
+  def test_instance_attribute(self):
+    attr = classes.Module(self.ctx, 'os').get_attribute('name')
+    self.assertIsInstance(attr, classes.FrozenInstance)
+    self.assertEqual(attr.cls.name, 'str')
+
+  def test_class_attribute(self):
+    attr = classes.Module(self.ctx, 'os').get_attribute('__name__')
+    self.assertIsInstance(attr, classes.FrozenInstance)
+    self.assertEqual(attr.cls.name, 'str')
+
+  def test_submodule(self):
+    attr = classes.Module(self.ctx, 'os').get_attribute('path')
+    self.assertIsInstance(attr, classes.Module)
+    self.assertEqual(attr.name, 'os.path')
+
+
 if __name__ == '__main__':
   unittest.main()
diff --git a/pytype/rewrite/convert.py b/pytype/rewrite/convert.py
@@ -107,3 +107,9 @@ def _pytd_type_to_value(self, typ: pytd.Type) -> abstract.BaseValue:
           f'Abstract conversion not yet implemented for {typ}')
     else:
       raise ValueError(f'Cannot convert {typ} to an abstract value')
+
+  def pytd_alias_to_value(self, alias: pytd.Alias) -> abstract.BaseValue:
+    if isinstance(alias.type, pytd.Module):
+      return abstract.Module(self._ctx, alias.type.module_name)
+    raise NotImplementedError(
+        f'Abstract conversion not yet implemented for {alias}')
diff --git a/pytype/rewrite/convert_test.py b/pytype/rewrite/convert_test.py
@@ -92,5 +92,14 @@ class D: ...
     self.assertEqual(nested_class.name, 'D')
 
 
+class PytdAliasToValueTest(ConverterTestBase):
+
+  def test_alias(self):
+    alias = self.build_pytd('import os.path', name='os.path')
+    module = self.conv.pytd_alias_to_value(alias)
+    self.assertIsInstance(module, abstract.Module)
+    self.assertEqual(module.name, 'os.path')
+
+
 if __name__ == '__main__':
   unittest.main()
diff --git a/pytype/rewrite/frame.py b/pytype/rewrite/frame.py
@@ -533,8 +533,17 @@ def byte_LOAD_METHOD(self, opcode):
 
   def byte_IMPORT_NAME(self, opcode):
     full_name = opcode.argval
-    unused_level_var, unused_fromlist = self._stack.popn(2)
-    module = abstract.Module(self._ctx, full_name)
+    unused_level_var, fromlist = self._stack.popn(2)
+    # The IMPORT_NAME for an "import a.b.c" will push the module "a".
+    # However, for "from a.b.c import Foo" it'll push the module "a.b.c". Those
+    # two cases are distinguished by whether fromlist is None or not.
+    try:
+      abstract.get_atomic_constant(fromlist, None)
+    except ValueError:
+      module_name = full_name
+    else:
+      module_name = full_name.split('.', 1)[0]  # "a.b.c" -> "a"
+    module = abstract.Module(self._ctx, module_name)
     return self._stack.push(module.to_variable())
 
   # ---------------------------------------------------------------
diff --git a/pytype/rewrite/load_abstract.py b/pytype/rewrite/load_abstract.py
@@ -60,6 +60,8 @@ def _load_pytd_node(self, pytd_node: pytd.Node) -> abstract.BaseValue:
     elif isinstance(pytd_node, pytd.Constant):
       typ = self._ctx.abstract_converter.pytd_type_to_value(pytd_node.type)
       return typ.instantiate()
+    elif isinstance(pytd_node, pytd.Alias):
+      return self._ctx.abstract_converter.pytd_alias_to_value(pytd_node)
     else:
       raise NotImplementedError(f'I do not know how to load {pytd_node}')
 
diff --git a/pytype/rewrite/load_abstract_test.py b/pytype/rewrite/load_abstract_test.py
@@ -50,6 +50,11 @@ def test_stdlib(self):
     self.assertIsInstance(name, abstract.FrozenInstance)
     self.assertEqual(name.cls.name, 'str')
 
+  def test_submodule(self):
+    submodule = self.ctx.abstract_loader.load_value('os', 'path')
+    self.assertIsInstance(submodule, abstract.Module)
+    self.assertEqual(submodule.name, 'os.path')
+
 
 class LoadRawTypeTest(test_utils.ContextfulTestBase):
 
diff --git a/pytype/rewrite/tests/test_basic.py b/pytype/rewrite/tests/test_basic.py
@@ -3,7 +3,14 @@
 from pytype.tests import test_base
 
 
-class BasicTest(test_base.BaseTest):
+class RewriteTest(test_base.BaseTest):
+
+  def setUp(self):
+    super().setUp()
+    self.options.tweak(use_rewrite=True)
+
+
+class BasicTest(RewriteTest):
   """Basic functional tests."""
 
   def setUp(self):
@@ -90,6 +97,10 @@ def __init__(self) -> None: ...
         def f(self) -> int: ...
     """)
 
+
+class ImportsTest(RewriteTest):
+  """Import tests."""
+
   def test_import(self):
     self.Check("""
       import os
@@ -102,6 +113,38 @@ def test_builtins(self):
       assert_type(__builtins__.int, "Type[int]")
     """)
 
+  def test_dotted_import(self):
+    self.Check("""
+      import os.path
+      assert_type(os.path, "module")
+    """)
+
+  @test_base.skip('Not yet implemented')
+  def test_from_import(self):
+    self.Check("""
+      from os import name, path
+      assert_type(name, "str")
+      assert_type(path, "module")
+    """)
+
+  @test_base.skip('Not yet implemented')
+  def test_errors(self):
+    self.CheckWithErrors("""
+      import nonsense  # import-error
+      import os.nonsense  # import-error
+      from os import nonsense  # import-error
+    """)
+
+  def test_aliases(self):
+    self.Check("""
+      import os as so
+      assert_type(so.name, "str")
+
+      # Not yet implemented:
+      # import os.path as path1
+      # from os import path as path2
+    """)
+
 
 if __name__ == '__main__':
   test_base.main()
