Commits

Boris FELD committed 380b4e6

Working version of assert hook

  • Participants
  • Parent commits 05ab526

Comments (0)

Files changed (6)

+import sys
+import imp
+import os.path
+from ast import *
+
+def new_assert(assertion):
+    print "Called with", assertion
+
+def assert_path_choose(path):
+    return True
+
+class RewriteAssert(NodeTransformer):
+    def __init__(self, content, filename):
+        self.content = content
+        self.filename = filename
+
+    def visit_Module(self, module_node):
+        # Need to import and call ast_report.register_module().
+        # Find where I can insert them.
+        body_future = []
+        body_rest = []
+        for node in module_node.body:
+            node = self.visit(node)
+            if (not body_rest and isinstance(node, ImportFrom) and
+                    node.module == "__future__"):
+                body_future.append(node)
+            else:
+                body_rest.append(node)
+
+        # It's easier to let Python convert the code to an AST
+        import_line = parse("from assert_hook import new_assert").body[0]
+
+        # Assign a reasonable seeming line number.
+        lineno = 1
+        if body_future:
+            lineno = body_future[0].lineno
+        for new_node in (import_line,):
+            new_node.col_offset = 1
+            new_node.lineno = lineno
+
+        new_body = body_future + [import_line] + body_rest
+        return Module(body=new_body)
+
+    def visit_Assert(self, node):
+        base = parse('new_assert("")').body[0]
+
+        assertion = self.content[node.lineno - 1].strip()[7:]
+        base.value.args[0].s = assertion
+        return base
+
+class ImportHook(object):
+
+    def find_module(self, module_name, package_path):
+        if assert_path_choose(module_name):
+            self.module_name = module_name
+            self.package_path = package_path[0] if package_path else ''
+            return self
+        return
+
+    def load_module(self, fullname):
+        if fullname in sys.modules:
+            return sys.modules[fullname]
+
+        fullpath = os.path.join(self.package_path, fullname.split('.')[-1] + '.py')
+        with open(fullpath, 'r') as source_code:
+            content = source_code.readlines()
+
+        code = parse(''.join(content), fullname)
+        code = RewriteAssert(content, fullname).visit(code)
+
+        mod = sys.modules.setdefault(fullname, imp.new_module(fullname))
+        mod.__file__ = "%s" % fullpath
+        mod.__loader__ = self
+
+        exec compile(code, fullpath, 'exec') in mod.__dict__
+        return mod
+
+sys.meta_path.append(ImportHook())
-import ast
-
-def new_assert(assertion):
-    print("Called with", ast.dump(assertion))
+def d():
+    return 1
 
 def test():
-    assert 1 == 3
+    assert d() == 2
 
 test()
+import assert_hook
+
+def choose_path(module_name):
+    return module_name.split('.')[-1].startswith('test')
+
+assert_hook.assert_path_choose = choose_path
+
+import test
 from ast import *
 from unparse import Unparser
 
+import assert_hook
+
 def new_assert(assertion):
-    print("Called with", dump(assertion))
+    print("Called with", assertion)
+    return
+
+assert_hook.new_assert = new_assert
 
 class RewriteAssert(NodeTransformer):
+    def __init__(self, content, filename):
+        self.content = content
+        self.filename = filename
+
+    def visit_Module(self, module_node):
+        # Need to import and call ast_report.register_module().
+        # Find where I can insert them.
+        body_future = []
+        body_rest = []
+        for node in module_node.body:
+            node = self.visit(node)
+            if (not body_rest and isinstance(node, ImportFrom) and
+                    node.module == "__future__"):
+                body_future.append(node)
+            else:
+                body_rest.append(node)
+
+        # It's easier to let Python convert the code to an AST
+        import_line = parse("from assert_hook import new_assert").body[0]
+
+        # Assign a reasonable seeming line number.
+        lineno = 1
+        if body_future:
+            lineno = body_future[0].lineno
+        for new_node in (import_line,):
+            new_node.col_offset = 1
+            new_node.lineno = lineno
+
+        new_body = body_future + [import_line] + body_rest
+        return Module(body=new_body)
+
     def visit_Assert(self, node):
-        base = parse('new_assert(ast.parse(""))')
-        stringio = StringIO.StringIO()
-        Unparser(node.test, stringio)
-        base.body[0].value.args[0].args[0].s = stringio.getvalue()
-        base = base.body[0]
+        base = parse('new_assert("")').body[0]
+
+        assertion = self.content[node.lineno - 1].strip()[7:]
+        base.value.args[0].s = assertion
         return base
 
 class ImportHook(object):
 
-    package_path = ''
-
     def find_module(self, module_name, package_path):
         if module_name.split('.')[-1].startswith('test'):
             self.module_name = module_name
-            self.package_path = package_path[0]
+            self.package_path = package_path[0] if package_path else ''
             return self
         return
 
 
         fullpath = os.path.join(self.package_path, fullname.split('.')[-1] + '.py')
         with open(fullpath, 'r') as source_code:
-            code = parse(source_code.read(), fullname)
+            content = source_code.readlines()
 
-        RewriteAssert().visit(code)
-        code = fix_missing_locations(code)
+        code = parse(''.join(content), fullname)
+
+        code = RewriteAssert(content, fullname).visit(code)
 
         mod = sys.modules.setdefault(fullname, imp.new_module(fullname))
         mod.__file__ = "%s" % fullpath
         mod.__loader__ = self
+
         exec compile(code, fullpath, 'exec') in mod.__dict__
         return mod
 
 sys.meta_path.append(ImportHook())
 
+import test
 from example import test
+import sys
+
 def test():
     assert 1 == 3
 

unparse.py

-"Usage: unparse.py <path to source file>"
-import sys
-import ast
-import cStringIO
-import os
-
-# Large float and imaginary literals get turned into infinities in the AST.
-# We unparse those infinities to INFSTR.
-INFSTR = "1e" + repr(sys.float_info.max_10_exp + 1)
-
-def interleave(inter, f, seq):
-    """Call f on each item in seq, calling inter() in between.
-    """
-    seq = iter(seq)
-    try:
-        f(next(seq))
-    except StopIteration:
-        pass
-    else:
-        for x in seq:
-            inter()
-            f(x)
-
-class Unparser:
-    """Methods in this class recursively traverse an AST and
-    output source code for the abstract syntax; original formatting
-    is disregarded. """
-
-    def __init__(self, tree, file = sys.stdout):
-        """Unparser(tree, file=sys.stdout) -> None.
-         Print the source for tree to file."""
-        self.f = file
-        self.future_imports = []
-        self._indent = 0
-        self.dispatch(tree)
-        self.f.write("")
-        self.f.flush()
-
-    def fill(self, text = ""):
-        "Indent a piece of text, according to the current indentation level"
-        self.f.write("\n"+"    "*self._indent + text)
-
-    def write(self, text):
-        "Append a piece of text to the current line."
-        self.f.write(text)
-
-    def enter(self):
-        "Print ':', and increase the indentation."
-        self.write(":")
-        self._indent += 1
-
-    def leave(self):
-        "Decrease the indentation level."
-        self._indent -= 1
-
-    def dispatch(self, tree):
-        "Dispatcher function, dispatching tree type T to method _T."
-        if isinstance(tree, list):
-            for t in tree:
-                self.dispatch(t)
-            return
-        meth = getattr(self, "_"+tree.__class__.__name__)
-        meth(tree)
-
-
-    ############### Unparsing methods ######################
-    # There should be one method per concrete grammar type #
-    # Constructors should be grouped by sum type. Ideally, #
-    # this would follow the order in the grammar, but      #
-    # currently doesn't.                                   #
-    ########################################################
-
-    def _Module(self, tree):
-        for stmt in tree.body:
-            self.dispatch(stmt)
-
-    # stmt
-    def _Expr(self, tree):
-        self.fill()
-        self.dispatch(tree.value)
-
-    def _Import(self, t):
-        self.fill("import ")
-        interleave(lambda: self.write(", "), self.dispatch, t.names)
-
-    def _ImportFrom(self, t):
-        # A from __future__ import may affect unparsing, so record it.
-        if t.module and t.module == '__future__':
-            self.future_imports.extend(n.name for n in t.names)
-
-        self.fill("from ")
-        self.write("." * t.level)
-        if t.module:
-            self.write(t.module)
-        self.write(" import ")
-        interleave(lambda: self.write(", "), self.dispatch, t.names)
-
-    def _Assign(self, t):
-        self.fill()
-        for target in t.targets:
-            self.dispatch(target)
-            self.write(" = ")
-        self.dispatch(t.value)
-
-    def _AugAssign(self, t):
-        self.fill()
-        self.dispatch(t.target)
-        self.write(" "+self.binop[t.op.__class__.__name__]+"= ")
-        self.dispatch(t.value)
-
-    def _Return(self, t):
-        self.fill("return")
-        if t.value:
-            self.write(" ")
-            self.dispatch(t.value)
-
-    def _Pass(self, t):
-        self.fill("pass")
-
-    def _Break(self, t):
-        self.fill("break")
-
-    def _Continue(self, t):
-        self.fill("continue")
-
-    def _Delete(self, t):
-        self.fill("del ")
-        interleave(lambda: self.write(", "), self.dispatch, t.targets)
-
-    def _Assert(self, t):
-        self.fill("assert ")
-        self.dispatch(t.test)
-        if t.msg:
-            self.write(", ")
-            self.dispatch(t.msg)
-
-    def _Exec(self, t):
-        self.fill("exec ")
-        self.dispatch(t.body)
-        if t.globals:
-            self.write(" in ")
-            self.dispatch(t.globals)
-        if t.locals:
-            self.write(", ")
-            self.dispatch(t.locals)
-
-    def _Print(self, t):
-        self.fill("print ")
-        do_comma = False
-        if t.dest:
-            self.write(">>")
-            self.dispatch(t.dest)
-            do_comma = True
-        for e in t.values:
-            if do_comma:self.write(", ")
-            else:do_comma=True
-            self.dispatch(e)
-        if not t.nl:
-            self.write(",")
-
-    def _Global(self, t):
-        self.fill("global ")
-        interleave(lambda: self.write(", "), self.write, t.names)
-
-    def _Yield(self, t):
-        self.write("(")
-        self.write("yield")
-        if t.value:
-            self.write(" ")
-            self.dispatch(t.value)
-        self.write(")")
-
-    def _Raise(self, t):
-        self.fill('raise ')
-        if t.type:
-            self.dispatch(t.type)
-        if t.inst:
-            self.write(", ")
-            self.dispatch(t.inst)
-        if t.tback:
-            self.write(", ")
-            self.dispatch(t.tback)
-
-    def _TryExcept(self, t):
-        self.fill("try")
-        self.enter()
-        self.dispatch(t.body)
-        self.leave()
-
-        for ex in t.handlers:
-            self.dispatch(ex)
-        if t.orelse:
-            self.fill("else")
-            self.enter()
-            self.dispatch(t.orelse)
-            self.leave()
-
-    def _TryFinally(self, t):
-        if len(t.body) == 1 and isinstance(t.body[0], ast.TryExcept):
-            # try-except-finally
-            self.dispatch(t.body)
-        else:
-            self.fill("try")
-            self.enter()
-            self.dispatch(t.body)
-            self.leave()
-
-        self.fill("finally")
-        self.enter()
-        self.dispatch(t.finalbody)
-        self.leave()
-
-    def _ExceptHandler(self, t):
-        self.fill("except")
-        if t.type:
-            self.write(" ")
-            self.dispatch(t.type)
-        if t.name:
-            self.write(" as ")
-            self.dispatch(t.name)
-        self.enter()
-        self.dispatch(t.body)
-        self.leave()
-
-    def _ClassDef(self, t):
-        self.write("\n")
-        for deco in t.decorator_list:
-            self.fill("@")
-            self.dispatch(deco)
-        self.fill("class "+t.name)
-        if t.bases:
-            self.write("(")
-            for a in t.bases:
-                self.dispatch(a)
-                self.write(", ")
-            self.write(")")
-        self.enter()
-        self.dispatch(t.body)
-        self.leave()
-
-    def _FunctionDef(self, t):
-        self.write("\n")
-        for deco in t.decorator_list:
-            self.fill("@")
-            self.dispatch(deco)
-        self.fill("def "+t.name + "(")
-        self.dispatch(t.args)
-        self.write(")")
-        self.enter()
-        self.dispatch(t.body)
-        self.leave()
-
-    def _For(self, t):
-        self.fill("for ")
-        self.dispatch(t.target)
-        self.write(" in ")
-        self.dispatch(t.iter)
-        self.enter()
-        self.dispatch(t.body)
-        self.leave()
-        if t.orelse:
-            self.fill("else")
-            self.enter()
-            self.dispatch(t.orelse)
-            self.leave()
-
-    def _If(self, t):
-        self.fill("if ")
-        self.dispatch(t.test)
-        self.enter()
-        self.dispatch(t.body)
-        self.leave()
-        # collapse nested ifs into equivalent elifs.
-        while (t.orelse and len(t.orelse) == 1 and
-               isinstance(t.orelse[0], ast.If)):
-            t = t.orelse[0]
-            self.fill("elif ")
-            self.dispatch(t.test)
-            self.enter()
-            self.dispatch(t.body)
-            self.leave()
-        # final else
-        if t.orelse:
-            self.fill("else")
-            self.enter()
-            self.dispatch(t.orelse)
-            self.leave()
-
-    def _While(self, t):
-        self.fill("while ")
-        self.dispatch(t.test)
-        self.enter()
-        self.dispatch(t.body)
-        self.leave()
-        if t.orelse:
-            self.fill("else")
-            self.enter()
-            self.dispatch(t.orelse)
-            self.leave()
-
-    def _With(self, t):
-        self.fill("with ")
-        self.dispatch(t.context_expr)
-        if t.optional_vars:
-            self.write(" as ")
-            self.dispatch(t.optional_vars)
-        self.enter()
-        self.dispatch(t.body)
-        self.leave()
-
-    # expr
-    def _Str(self, tree):
-        # if from __future__ import unicode_literals is in effect,
-        # then we want to output string literals using a 'b' prefix
-        # and unicode literals with no prefix.
-        if "unicode_literals" not in self.future_imports:
-            self.write(repr(tree.s))
-        elif isinstance(tree.s, str):
-            self.write("b" + repr(tree.s))
-        elif isinstance(tree.s, unicode):
-            self.write(repr(tree.s).lstrip("u"))
-        else:
-            assert False, "shouldn't get here"
-
-    def _Name(self, t):
-        self.write(t.id)
-
-    def _Repr(self, t):
-        self.write("`")
-        self.dispatch(t.value)
-        self.write("`")
-
-    def _Num(self, t):
-        repr_n = repr(t.n)
-        # Parenthesize negative numbers, to avoid turning (-1)**2 into -1**2.
-        if repr_n.startswith("-"):
-            self.write("(")
-        # Substitute overflowing decimal literal for AST infinities.
-        self.write(repr_n.replace("inf", INFSTR))
-        if repr_n.startswith("-"):
-            self.write(")")
-
-    def _List(self, t):
-        self.write("[")
-        interleave(lambda: self.write(", "), self.dispatch, t.elts)
-        self.write("]")
-
-    def _ListComp(self, t):
-        self.write("[")
-        self.dispatch(t.elt)
-        for gen in t.generators:
-            self.dispatch(gen)
-        self.write("]")
-
-    def _GeneratorExp(self, t):
-        self.write("(")
-        self.dispatch(t.elt)
-        for gen in t.generators:
-            self.dispatch(gen)
-        self.write(")")
-
-    def _SetComp(self, t):
-        self.write("{")
-        self.dispatch(t.elt)
-        for gen in t.generators:
-            self.dispatch(gen)
-        self.write("}")
-
-    def _DictComp(self, t):
-        self.write("{")
-        self.dispatch(t.key)
-        self.write(": ")
-        self.dispatch(t.value)
-        for gen in t.generators:
-            self.dispatch(gen)
-        self.write("}")
-
-    def _comprehension(self, t):
-        self.write(" for ")
-        self.dispatch(t.target)
-        self.write(" in ")
-        self.dispatch(t.iter)
-        for if_clause in t.ifs:
-            self.write(" if ")
-            self.dispatch(if_clause)
-
-    def _IfExp(self, t):
-        self.write("(")
-        self.dispatch(t.body)
-        self.write(" if ")
-        self.dispatch(t.test)
-        self.write(" else ")
-        self.dispatch(t.orelse)
-        self.write(")")
-
-    def _Set(self, t):
-        assert(t.elts) # should be at least one element
-        self.write("{")
-        interleave(lambda: self.write(", "), self.dispatch, t.elts)
-        self.write("}")
-
-    def _Dict(self, t):
-        self.write("{")
-        def write_pair(pair):
-            (k, v) = pair
-            self.dispatch(k)
-            self.write(": ")
-            self.dispatch(v)
-        interleave(lambda: self.write(", "), write_pair, zip(t.keys, t.values))
-        self.write("}")
-
-    def _Tuple(self, t):
-        self.write("(")
-        if len(t.elts) == 1:
-            (elt,) = t.elts
-            self.dispatch(elt)
-            self.write(",")
-        else:
-            interleave(lambda: self.write(", "), self.dispatch, t.elts)
-        self.write(")")
-
-    unop = {"Invert":"~", "Not": "not", "UAdd":"+", "USub":"-"}
-    def _UnaryOp(self, t):
-        self.write("(")
-        self.write(self.unop[t.op.__class__.__name__])
-        self.write(" ")
-        # If we're applying unary minus to a number, parenthesize the number.
-        # This is necessary: -2147483648 is different from -(2147483648) on
-        # a 32-bit machine (the first is an int, the second a long), and
-        # -7j is different from -(7j).  (The first has real part 0.0, the second
-        # has real part -0.0.)
-        if isinstance(t.op, ast.USub) and isinstance(t.operand, ast.Num):
-            self.write("(")
-            self.dispatch(t.operand)
-            self.write(")")
-        else:
-            self.dispatch(t.operand)
-        self.write(")")
-
-    binop = { "Add":"+", "Sub":"-", "Mult":"*", "Div":"/", "Mod":"%",
-                    "LShift":"<<", "RShift":">>", "BitOr":"|", "BitXor":"^", "BitAnd":"&",
-                    "FloorDiv":"//", "Pow": "**"}
-    def _BinOp(self, t):
-        self.write("(")
-        self.dispatch(t.left)
-        self.write(" " + self.binop[t.op.__class__.__name__] + " ")
-        self.dispatch(t.right)
-        self.write(")")
-
-    cmpops = {"Eq":"==", "NotEq":"!=", "Lt":"<", "LtE":"<=", "Gt":">", "GtE":">=",
-                        "Is":"is", "IsNot":"is not", "In":"in", "NotIn":"not in"}
-    def _Compare(self, t):
-        self.write("(")
-        self.dispatch(t.left)
-        for o, e in zip(t.ops, t.comparators):
-            self.write(" " + self.cmpops[o.__class__.__name__] + " ")
-            self.dispatch(e)
-        self.write(")")
-
-    boolops = {ast.And: 'and', ast.Or: 'or'}
-    def _BoolOp(self, t):
-        self.write("(")
-        s = " %s " % self.boolops[t.op.__class__]
-        interleave(lambda: self.write(s), self.dispatch, t.values)
-        self.write(")")
-
-    def _Attribute(self,t):
-        self.dispatch(t.value)
-        # Special case: 3.__abs__() is a syntax error, so if t.value
-        # is an integer literal then we need to either parenthesize
-        # it or add an extra space to get 3 .__abs__().
-        if isinstance(t.value, ast.Num) and isinstance(t.value.n, int):
-            self.write(" ")
-        self.write(".")
-        self.write(t.attr)
-
-    def _Call(self, t):
-        self.dispatch(t.func)
-        self.write("(")
-        comma = False
-        for e in t.args:
-            if comma: self.write(", ")
-            else: comma = True
-            self.dispatch(e)
-        for e in t.keywords:
-            if comma: self.write(", ")
-            else: comma = True
-            self.dispatch(e)
-        if t.starargs:
-            if comma: self.write(", ")
-            else: comma = True
-            self.write("*")
-            self.dispatch(t.starargs)
-        if t.kwargs:
-            if comma: self.write(", ")
-            else: comma = True
-            self.write("**")
-            self.dispatch(t.kwargs)
-        self.write(")")
-
-    def _Subscript(self, t):
-        self.dispatch(t.value)
-        self.write("[")
-        self.dispatch(t.slice)
-        self.write("]")
-
-    # slice
-    def _Ellipsis(self, t):
-        self.write("...")
-
-    def _Index(self, t):
-        self.dispatch(t.value)
-
-    def _Slice(self, t):
-        if t.lower:
-            self.dispatch(t.lower)
-        self.write(":")
-        if t.upper:
-            self.dispatch(t.upper)
-        if t.step:
-            self.write(":")
-            self.dispatch(t.step)
-
-    def _ExtSlice(self, t):
-        interleave(lambda: self.write(', '), self.dispatch, t.dims)
-
-    # others
-    def _arguments(self, t):
-        first = True
-        # normal arguments
-        defaults = [None] * (len(t.args) - len(t.defaults)) + t.defaults
-        for a,d in zip(t.args, defaults):
-            if first:first = False
-            else: self.write(", ")
-            self.dispatch(a),
-            if d:
-                self.write("=")
-                self.dispatch(d)
-
-        # varargs
-        if t.vararg:
-            if first:first = False
-            else: self.write(", ")
-            self.write("*")
-            self.write(t.vararg)
-
-        # kwargs
-        if t.kwarg:
-            if first:first = False
-            else: self.write(", ")
-            self.write("**"+t.kwarg)
-
-    def _keyword(self, t):
-        self.write(t.arg)
-        self.write("=")
-        self.dispatch(t.value)
-
-    def _Lambda(self, t):
-        self.write("(")
-        self.write("lambda ")
-        self.dispatch(t.args)
-        self.write(": ")
-        self.dispatch(t.body)
-        self.write(")")
-
-    def _alias(self, t):
-        self.write(t.name)
-        if t.asname:
-            self.write(" as "+t.asname)
-
-def roundtrip(filename, output=sys.stdout):
-    with open(filename, "r") as pyfile:
-        source = pyfile.read()
-    tree = compile(source, filename, "exec", ast.PyCF_ONLY_AST)
-    Unparser(tree, output)
-
-
-
-def testdir(a):
-    try:
-        names = [n for n in os.listdir(a) if n.endswith('.py')]
-    except OSError:
-        sys.stderr.write("Directory not readable: %s" % a)
-    else:
-        for n in names:
-            fullname = os.path.join(a, n)
-            if os.path.isfile(fullname):
-                output = cStringIO.StringIO()
-                print 'Testing %s' % fullname
-                try:
-                    roundtrip(fullname, output)
-                except Exception as e:
-                    print '  Failed to compile, exception is %s' % repr(e)
-            elif os.path.isdir(fullname):
-                testdir(fullname)
-
-def main(args):
-    if args[0] == '--testdir':
-        for a in args[1:]:
-            testdir(a)
-    else:
-        for a in args:
-            roundtrip(a)
-
-if __name__=='__main__':
-    main(sys.argv[1:])