Add support for absolute/relative imports and if/else expressions:

 - regenerate ast.py
 - add future flags for absolute-import and with-statement so they
   (hopefully) properly get set in code-object flags
 - try out if/else expressions in actual code for the hell of it.

Seems to generate the same kind of bytecode as the normal compiler.

5 files changed, 82 insertions, 19 deletions

diff --git a/Lib/compiler/ast.py b/Lib/compiler/ast.py
index 6e3e182..7d1c18f 100644
--- a/Lib/compiler/ast.py
+++ b/Lib/compiler/ast.py
@@ -524,19 +524,20 @@ class For(Node):
         return "For(%s, %s, %s, %s)" % (repr(self.assign), repr(self.list), repr(self.body), repr(self.else_))
 
 class From(Node):
-    def __init__(self, modname, names, lineno=None):
+    def __init__(self, modname, names, level, lineno=None):
         self.modname = modname
         self.names = names
+        self.level = level
         self.lineno = lineno
 
     def getChildren(self):
-        return self.modname, self.names
+        return self.modname, self.names, self.level
 
     def getChildNodes(self):
         return ()
 
     def __repr__(self):
-        return "From(%s, %s)" % (repr(self.modname), repr(self.names))
+        return "From(%s, %s, %s)" % (repr(self.modname), repr(self.names), repr(self.level))
 
 class Function(Node):
     def __init__(self, decorators, name, argnames, defaults, flags, doc, code, lineno=None):
@@ -553,7 +554,7 @@ class Function(Node):
             self.varargs = 1
         if flags & CO_VARKEYWORDS:
             self.kwargs = 1
-
+    
 
 
     def getChildren(self):
@@ -584,7 +585,7 @@ class GenExpr(Node):
         self.lineno = lineno
         self.argnames = ['[outmost-iterable]']
         self.varargs = self.kwargs = None
-
+    
 
 
     def getChildren(self):
@@ -708,6 +709,22 @@ class If(Node):
     def __repr__(self):
         return "If(%s, %s)" % (repr(self.tests), repr(self.else_))
 
+class IfExp(Node):
+    def __init__(self, test, then, else_, lineno=None):
+        self.test = test
+        self.then = then
+        self.else_ = else_
+        self.lineno = lineno
+
+    def getChildren(self):
+        return self.test, self.then, self.else_
+
+    def getChildNodes(self):
+        return self.test, self.then, self.else_
+
+    def __repr__(self):
+        return "IfExp(%s, %s, %s)" % (repr(self.test), repr(self.then), repr(self.else_))
+
 class Import(Node):
     def __init__(self, names, lineno=None):
         self.names = names
@@ -763,7 +780,7 @@ class Lambda(Node):
             self.varargs = 1
         if flags & CO_VARKEYWORDS:
             self.kwargs = 1
-
+    
 
 
     def getChildren(self):
diff --git a/Lib/compiler/consts.py b/Lib/compiler/consts.py
index a6cf559..d760982 100644
--- a/Lib/compiler/consts.py
+++ b/Lib/compiler/consts.py
@@ -17,3 +17,5 @@ CO_NESTED = 0x0010
 CO_GENERATOR = 0x0020
 CO_GENERATOR_ALLOWED = 0x1000
 CO_FUTURE_DIVISION = 0x2000
+CO_FUTURE_ABSIMPORT = 0x4000
+CO_FUTURE_WITH_STATEMENT = 0x8000
diff --git a/Lib/compiler/pyassem.py b/Lib/compiler/pyassem.py
index 22a19c7..5932ffe 100644
--- a/Lib/compiler/pyassem.py
+++ b/Lib/compiler/pyassem.py
@@ -771,7 +771,7 @@ class StackDepthTracker:
         'COMPARE_OP': -1,
         'STORE_FAST': -1,
         'IMPORT_STAR': -1,
-        'IMPORT_NAME': 0,
+        'IMPORT_NAME': -1,
         'IMPORT_FROM': 1,
         'LOAD_ATTR': 0, # unlike other loads
         # close enough...
diff --git a/Lib/compiler/pycodegen.py b/Lib/compiler/pycodegen.py
index 441ccf3..9959b13 100644
--- a/Lib/compiler/pycodegen.py
+++ b/Lib/compiler/pycodegen.py
@@ -8,8 +8,9 @@ from cStringIO import StringIO
 from compiler import ast, parse, walk, syntax
 from compiler import pyassem, misc, future, symbols
 from compiler.consts import SC_LOCAL, SC_GLOBAL, SC_FREE, SC_CELL
-from compiler.consts import CO_VARARGS, CO_VARKEYWORDS, CO_NEWLOCALS,\
-     CO_NESTED, CO_GENERATOR, CO_GENERATOR_ALLOWED, CO_FUTURE_DIVISION
+from compiler.consts import (CO_VARARGS, CO_VARKEYWORDS, CO_NEWLOCALS,
+     CO_NESTED, CO_GENERATOR, CO_GENERATOR_ALLOWED, CO_FUTURE_DIVISION,
+     CO_FUTURE_ABSIMPORT, CO_FUTURE_WITH_STATEMENT)
 from compiler.pyassem import TupleArg
 
 # XXX The version-specific code can go, since this code only works with 2.x.
@@ -215,6 +216,10 @@ class CodeGenerator:
                 self._div_op = "BINARY_TRUE_DIVIDE"
             elif feature == "generators":
                 self.graph.setFlag(CO_GENERATOR_ALLOWED)
+            elif feature == "absolute_import":
+                self.graph.setFlag(CO_FUTURE_ABSIMPORT)
+            elif feature == "with_statement":
+                self.graph.setFlag(CO_FUTURE_WITH_STATEMENT)
 
     def initClass(self):
         """This method is called once for each class"""
@@ -543,6 +548,19 @@ class CodeGenerator:
     def visitOr(self, node):
         self.visitTest(node, 'JUMP_IF_TRUE')
 
+    def visitIfExp(self, node):
+        endblock = self.newBlock()
+        elseblock = self.newBlock()
+        self.visit(node.test)
+        self.emit('JUMP_IF_FALSE', elseblock)
+        self.emit('POP_TOP')
+        self.visit(node.then)
+        self.emit('JUMP_FORWARD', endblock)
+        self.nextBlock(elseblock)
+        self.emit('POP_TOP')
+        self.visit(node.else_)
+        self.nextBlock(endblock)
+
     def visitCompare(self, node):
         self.visit(node.expr)
         cleanup = self.newBlock()
@@ -875,8 +893,10 @@ class CodeGenerator:
 
     def visitImport(self, node):
         self.set_lineno(node)
+        level = 0 if "absolute_import" in self.futures else -1
         for name, alias in node.names:
             if VERSION > 1:
+                self.emit('LOAD_CONST', level)
                 self.emit('LOAD_CONST', None)
             self.emit('IMPORT_NAME', name)
             mod = name.split(".")[0]
@@ -888,8 +908,12 @@ class CodeGenerator:
 
     def visitFrom(self, node):
         self.set_lineno(node)
+        level = node.level
+        if level == 0 and "absolute_import" not in self.futures:
+            level = -1
         fromlist = map(lambda (name, alias): name, node.names)
         if VERSION > 1:
+            self.emit('LOAD_CONST', level)
             self.emit('LOAD_CONST', tuple(fromlist))
         self.emit('IMPORT_NAME', node.modname)
         for name, alias in node.names:
diff --git a/Lib/compiler/transformer.py b/Lib/compiler/transformer.py
index 800461c..cc91b4f 100644
--- a/Lib/compiler/transformer.py
+++ b/Lib/compiler/transformer.py
@@ -441,18 +441,25 @@ class Transformer:
                       lineno=nodelist[0][2])
 
     def import_from(self, nodelist):
-        # import_from: 'from' dotted_name 'import' ('*' |
+        # import_from: 'from' ('.'* dotted_name | '.') 'import' ('*' |
         #    '(' import_as_names ')' | import_as_names)
         assert nodelist[0][1] == 'from'
-        assert nodelist[1][0] == symbol.dotted_name
-        assert nodelist[2][1] == 'import'
-        fromname = self.com_dotted_name(nodelist[1])
-        if nodelist[3][0] == token.STAR:
-            return From(fromname, [('*', None)],
+        idx = 1
+        while nodelist[idx][1] == '.':
+            idx += 1
+        level = idx - 1
+        if nodelist[idx][0] == symbol.dotted_name:
+            fromname = self.com_dotted_name(nodelist[idx])
+            idx += 1
+        else:
+            fromname = ""
+        assert nodelist[idx][1] == 'import'
+        if nodelist[idx + 1][0] == token.STAR:
+            return From(fromname, [('*', None)], level,
                         lineno=nodelist[0][2])
         else:
-            node = nodelist[3 + (nodelist[3][0] == token.LPAR)]
-            return From(fromname, self.com_import_as_names(node),
+            node = nodelist[idx + 1 + (nodelist[idx + 1][0] == token.LPAR)]
+            return From(fromname, self.com_import_as_names(node), level,
                         lineno=nodelist[0][2])
 
     def global_stmt(self, nodelist):
@@ -575,12 +582,25 @@ class Transformer:
         return self.testlist(nodelist)
 
     def test(self, nodelist):
+        # or_test ['if' or_test 'else' test] | lambdef
+        if len(nodelist) == 1 and nodelist[0][0] == symbol.lambdef:
+            return self.lambdef(nodelist[0])
+        then = self.com_node(nodelist[0])
+        if len(nodelist) > 1:
+            assert len(nodelist) == 5
+            assert nodelist[1][1] == 'if'
+            assert nodelist[3][1] == 'else'
+            test = self.com_node(nodelist[2])
+            else_ = self.com_node(nodelist[4])
+            return IfExp(test, then, else_, lineno=nodelist[1][2])
+        return then
+
+    def or_test(self, nodelist):
         # and_test ('or' and_test)* | lambdef
         if len(nodelist) == 1 and nodelist[0][0] == symbol.lambdef:
             return self.lambdef(nodelist[0])
         return self.com_binary(Or, nodelist)
-    or_test = test
-    old_test = test
+    old_test = or_test
 
     def and_test(self, nodelist):
         # not_test ('and' not_test)*