Adding relative import support (#60)

* Changing the `all_imported_modules_in_file(self, file)` projection in build which gathers imports from file ASTs root nodes.
* Adding support for the relative import syntax in noderepr.py and output.py.
* Adding a relative counter int to the AST nodes and fixing up treetransform.py.
* Changing the `parse_import_from(self)` function to parse relative imports.

Author: mason-bially

mypy/build.py

@@ -16,7 +16,7 @@
 import sys
 from os.path import dirname, basename
 
-from typing import Undefined, Dict, List, Tuple, cast, Set
+from typing import Undefined, Dict, List, Tuple, cast, Set, Union
 
 from mypy.types import Type
 from mypy.nodes import MypyFile, Node, Import, ImportFrom, ImportAll
@@ -451,21 +451,34 @@ def all_imported_modules_in_file(self,
         Return list of tuples (module id, import line number) for all modules
         imported in file.
         """
+        def correct_rel_imp(imp: Union[ImportFrom, ImportAll]) -> str:
+            """Function to correct for relative imports."""
+            file_id = file.fullname()
+            rel = imp.relative
+            if rel == 0:
+                return imp.id
+            if os.path.basename(file.path) == '__init__.py':
+                rel -= 1
+            if rel != 0:
+                file_id = ".".join(file_id.split(".")[:-rel])
+            return file_id + "." + imp.id if imp.id else file_id
+
         res = List[Tuple[str, int]]()
         for imp in file.imports:
             if not imp.is_unreachable:
                 if isinstance(imp, Import):
                     for id, _ in imp.ids:
                         res.append((id, imp.line))
                 elif isinstance(imp, ImportFrom):
-                    res.append((imp.id, imp.line))
+                    cur_id = correct_rel_imp(imp)
+                    res.append((cur_id, imp.line))
                     # Also add any imported names that are submodules.
                     for name, __ in imp.names:
-                        sub_id = imp.id + '.' + name
+                        sub_id = cur_id + '.' + name
                         if self.is_module(sub_id):
                             res.append((sub_id, imp.line))
                 elif isinstance(imp, ImportAll):
-                    res.append((imp.id, imp.line))
+                    res.append((correct_rel_imp(imp), imp.line))
         return res
 
     def is_module(self, id: str) -> bool:

mypy/noderepr.py

@@ -34,16 +34,18 @@ def __init__(self, import_tok: Any, components: List[List[Token]],
 class ImportFromRepr:
     def __init__(self,
                  from_tok: Any,
+                 rel_toks: Any,
                  components: List[Token],
                  import_tok: Any,
                  lparen: Any,
                  names: List[Tuple[List[Token], Token]],
                  rparen: Any, br: Any) -> None:
         # Notes:
-        # - lparen and rparen may be empty
+        # - lparen, rparen, and rel_tok may be empty
         # - in each names tuple, the first item contains tokens for
         #   'name [as name]' and the second item is a comma or empty.
         self.from_tok = from_tok
+        self.rel_toks = rel_toks
         self.components = components
         self.import_tok = import_tok
         self.lparen = lparen

mypy/nodes.py

@@ -163,20 +163,22 @@ class ImportFrom(ImportBase):
 
     names = Undefined(List[Tuple[str, str]])  # Tuples (name, as name)
 
-    def __init__(self, id: str, names: List[Tuple[str, str]]) -> None:
+    def __init__(self, id: str, relative: int, names: List[Tuple[str, str]]) -> None:
         self.id = id
         self.names = names
-
+        self.relative = relative
+    
     def accept(self, visitor: NodeVisitor[T]) -> T:
         return visitor.visit_import_from(self)
 
 
 class ImportAll(ImportBase):
     """from m import *"""
 
-    def __init__(self, id: str) -> None:
+    def __init__(self, id: str, relative: int) -> None:
         self.id = id
-
+        self.relative = relative
+    
     def accept(self, visitor: NodeVisitor[T]) -> T:
         return visitor.visit_import_all(self)
 

mypy/output.py

@@ -55,6 +55,7 @@ def visit_import_all(self, o):
     def output_import_from_or_all(self, o):
         r = o.repr
         self.token(r.from_tok)
+        self.tokens(r.rel_toks)
         self.tokens(r.components)
         self.token(r.import_tok)
         self.token(r.lparen)

mypy/parse.py

@@ -170,7 +170,16 @@ def parse_import(self) -> Import:
 
     def parse_import_from(self) -> Node:
         from_tok = self.expect('from')
-        name, components = self.parse_qualified_name()
+        relative = 0
+        rel_toks = List[Token]()
+        while self.current_str() == ".":
+            rel_toks.append(self.expect('.'))
+            relative += 1
+        if self.current_str() == "import":
+            name = ""
+            components = List[Token]()
+        else:
+            name, components = self.parse_qualified_name()
         if name == self.custom_typing_module:
             name = 'typing'
         import_tok = self.expect('import')
@@ -180,7 +189,7 @@ def parse_import_from(self) -> Node:
         node = None  # type: ImportBase
         if self.current_str() == '*':
             name_toks.append(([self.skip()], none))
-            node = ImportAll(name)
+            node = ImportAll(name, relative)
         else:
             is_paren = self.current_str() == '('
             if is_paren:
@@ -206,12 +215,12 @@ def parse_import_from(self) -> Node:
             if is_paren:
                 rparen = self.expect(')')
             if node is None:
-                node = ImportFrom(name, targets)
+                node = ImportFrom(name, relative, targets)
         br = self.expect_break()
         self.imports.append(node)
         # TODO: Fix representation if there is a custom typing module import.
         self.set_repr(node, noderepr.ImportFromRepr(
-            from_tok, components, import_tok, lparen, name_toks, rparen, br))
+            from_tok, rel_toks, components, import_tok, lparen, name_toks, rparen, br))
         if name == '__future__':
             self.future_options.extend(target[0] for target in targets)
         return node

mypy/treetransform.py

@@ -60,10 +60,10 @@ def visit_import(self, node: Import) -> Node:
         return Import(node.ids[:])
 
     def visit_import_from(self, node: ImportFrom) -> Node:
-        return ImportFrom(node.id, node.names[:])
+        return ImportFrom(node.id, node.relative, node.names[:])
 
     def visit_import_all(self, node: ImportAll) -> Node:
-        return ImportAll(node.id)
+        return ImportAll(node.id, node.relative)
 
     def visit_func_def(self, node: FuncDef) -> FuncDef:
         # Note that a FuncDef must be transformed to a FuncDef.