Render leading dots for relative imports (PEP 328)

ASTNode.h

@@ -315,20 +315,23 @@ class ASTImport : public ASTNode {
 public:
     typedef std::list<PycRef<ASTStore>> list_t;
 
-    ASTImport(PycRef<ASTNode> name, PycRef<ASTNode> fromlist)
-        : ASTNode(NODE_IMPORT), m_name(std::move(name)), m_fromlist(std::move(fromlist)) { }
+    ASTImport(PycRef<ASTNode> name, PycRef<ASTNode> fromlist, int level = 0)
+        : ASTNode(NODE_IMPORT), m_name(std::move(name)),
+          m_fromlist(std::move(fromlist)), m_level(level) { }
 
     PycRef<ASTNode> name() const { return m_name; }
     list_t stores() const { return m_stores; }
     void add_store(PycRef<ASTStore> store) { m_stores.emplace_back(std::move(store)); }
 
     PycRef<ASTNode> fromlist() const { return m_fromlist; }
+    int level() const { return m_level; }
 
 private:
     PycRef<ASTNode> m_name;
     list_t m_stores;
 
     PycRef<ASTNode> m_fromlist;
+    int m_level;
 };
 
 

ASTree.cpp

@@ -1095,9 +1095,17 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
             } else {
                 PycRef<ASTNode> fromlist = stack.top();
                 stack.pop();
-                if (mod->verCompare(2, 5) >= 0)
-                    stack.pop();    // Level -- we don't care
-                stack.push(new ASTImport(new ASTName(code->getName(operand)), fromlist));
+                int level = 0;
+                if (mod->verCompare(2, 5) >= 0) {
+                    PycRef<ASTNode> levelnode = stack.top();
+                    stack.pop();
+                    if (levelnode.type() == ASTNode::NODE_OBJECT) {
+                        PycRef<PycObject> obj = levelnode.cast<ASTObject>()->object();
+                        if (obj->type() == PycObject::TYPE_INT)
+                            level = obj.cast<PycInt>()->value();
+                    }
+                }
+                stack.push(new ASTImport(new ASTName(code->getName(operand)), fromlist, level));
             }
             break;
         case Pyc::IMPORT_FROM_A:
@@ -3259,6 +3267,8 @@ void print_src(PycRef<ASTNode> node, PycModule* mod, std::ostream& pyc_output)
                 ASTImport::list_t stores = import->stores();
 
                 pyc_output << "from ";
+                for (int i = 0; i < import->level(); i++)
+                    pyc_output << ".";
                 if (import->name().type() == ASTNode::NODE_IMPORT)
                     print_src(import->name().cast<ASTImport>()->name(), mod, pyc_output);
                 else
@@ -3438,6 +3448,8 @@ void print_src(PycRef<ASTNode> node, PycModule* mod, std::ostream& pyc_output)
                     PycRef<PycObject> fromlist = import->fromlist().cast<ASTObject>()->object();
                     if (fromlist != Pyc_None) {
                         pyc_output << "from ";
+                        for (int i = 0; i < import->level(); i++)
+                            pyc_output << ".";
                         if (import->name().type() == ASTNode::NODE_IMPORT)
                             print_src(import->name().cast<ASTImport>()->name(), mod, pyc_output);
                         else

tests/input/relative_imports.py

@@ -0,0 +1,5 @@
+from . import foo
+from .. import bar
+from .mod import baz
+from ..pkg.submod import qux, quux
+from ...deep import thing as renamed

tests/tokenized/relative_imports.txt

@@ -0,0 +1,5 @@
+from . import foo <EOL>
+from . . import bar <EOL>
+from . mod import baz <EOL>
+from . . pkg . submod import qux , quux <EOL>
+from ... deep import thing as renamed <EOL>