m-labs · Jul 2, 2015 · Jul 2, 2015 · Jul 2, 2015 · Jul 2, 2015
diff --git a/artiq/compiler/__init__.py b/artiq/compiler/__init__.py
@@ -0,0 +1 @@
+from .module import Module
diff --git a/artiq/compiler/module.py b/artiq/compiler/module.py
@@ -0,0 +1,36 @@
+"""
+The :class:`Module` class encapsulates a single Python
+"""
+
+import os
+from pythonparser import source, diagnostic, parse_buffer
+from . import prelude, types, transforms
+
+class Module:
+    def __init__(self, source_buffer, engine=diagnostic.Engine(all_errors_are_fatal=True)):
+        asttyped_rewriter = transforms.ASTTypedRewriter(engine=engine)
+        inferencer = transforms.Inferencer(engine=engine)
+        int_monomorphizer = transforms.IntMonomorphizer(engine=engine)
+
+        parsetree, comments = parse_buffer(source_buffer, engine=engine)
+        typedtree = asttyped_rewriter.visit(parsetree)
+        inferencer.visit(typedtree)
+        int_monomorphizer.visit(typedtree)
+        inferencer.visit(typedtree)
+
+        self.name = os.path.basename(source_buffer.name)
+        self.globals = asttyped_rewriter.globals
+
+    @classmethod
+    def from_string(klass, source_string, name="input.py", first_line=1):
+        return klass(source.Buffer(source_string + "\n", name, first_line))
+
+    @classmethod
+    def from_filename(klass, filename):
+        with open(filename) as f:
+            return klass(source.Buffer(f.read(), filename, 1))
+
+    def __repr__(self):
+        printer = types.TypePrinter()
+        globals = ["%s: %s" % (var, printer.name(self.globals[var])) for var in self.globals]
+        return "<artiq.compiler.Module %s {\n  %s\n}>" % (repr(self.name), ",\n  ".join(globals))
diff --git a/artiq/compiler/testbench/__init__.py b/artiq/compiler/testbench/__init__.py
diff --git a/artiq/compiler/testbench/inferencer.py b/artiq/compiler/testbench/inferencer.py
@@ -0,0 +1,75 @@
+import sys, fileinput, os
+from pythonparser import source, diagnostic, algorithm, parse_buffer
+from .. import prelude, types
+from ..transforms import ASTTypedRewriter, Inferencer
+
+class Printer(algorithm.Visitor):
+    """
+    :class:`Printer` prints ``:`` and the node type after every typed node,
+    and ``->`` and the node type before the colon in a function definition.
+
+    In almost all cases (except function definition) this does not result
+    in valid Python syntax.
+
+    :ivar rewriter: (:class:`pythonparser.source.Rewriter`) rewriter instance
+    """
+
+    def __init__(self, buf):
+        self.rewriter = source.Rewriter(buf)
+        self.type_printer = types.TypePrinter()
+
+    def rewrite(self):
+        return self.rewriter.rewrite()
+
+    def visit_FunctionDefT(self, node):
+        super().generic_visit(node)
+
+        self.rewriter.insert_before(node.colon_loc,
+                                    "->{}".format(self.type_printer.name(node.return_type)))
+
+    def visit_ExceptHandlerT(self, node):
+        super().generic_visit(node)
+
+        if node.name_loc:
+            self.rewriter.insert_after(node.name_loc,
+                                        ":{}".format(self.type_printer.name(node.name_type)))
+
+    def generic_visit(self, node):
+        super().generic_visit(node)
+
+        if hasattr(node, "type"):
+            self.rewriter.insert_after(node.loc,
+                                       ":{}".format(self.type_printer.name(node.type)))
+
+def main():
+    if len(sys.argv) > 1 and sys.argv[1] == '+diag':
+        del sys.argv[1]
+        def process_diagnostic(diag):
+            print("\n".join(diag.render(only_line=True)))
+            if diag.level == 'fatal':
+                exit()
+    else:
+        def process_diagnostic(diag):
+            print("\n".join(diag.render()))
+            if diag.level in ('fatal', 'error'):
+                exit(1)
+
+    engine = diagnostic.Engine()
+    engine.process = process_diagnostic
+
+    buf = source.Buffer("".join(fileinput.input()).expandtabs(),
+                        os.path.basename(fileinput.filename()))
+    parsed, comments = parse_buffer(buf, engine=engine)
+    typed = ASTTypedRewriter(globals=prelude.globals(), engine=engine).visit(parsed)
+    Inferencer(engine=engine).visit(typed)
+
+    printer = Printer(buf)
+    printer.visit(typed)
+    for comment in comments:
+        if comment.text.find("CHECK") >= 0:
+            printer.rewriter.remove(comment.loc)
+    print(printer.rewrite().source)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/artiq/compiler/transforms/__init__.py b/artiq/compiler/transforms/__init__.py
@@ -0,0 +1,3 @@
+from .asttyped_rewriter import ASTTypedRewriter
+from .inferencer import Inferencer
+from .int_monomorphizer import IntMonomorphizer