From fd5cdb7d1e25832fc99cc3bcec7c220a4275f056 Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Thu, 2 Mar 2017 15:06:21 +0000
Subject: [PATCH] compiler.transforms: implement a typedtree printer.

---
 artiq/compiler/transforms/__init__.py         |  1 +
 .../compiler/transforms/typedtree_printer.py  | 89 +++++++++++++++++++
 2 files changed, 90 insertions(+)
 create mode 100644 artiq/compiler/transforms/typedtree_printer.py

diff --git a/artiq/compiler/transforms/__init__.py b/artiq/compiler/transforms/__init__.py
index 665fd3ea1..e4275be36 100644
--- a/artiq/compiler/transforms/__init__.py
+++ b/artiq/compiler/transforms/__init__.py
@@ -6,3 +6,4 @@ from .artiq_ir_generator import ARTIQIRGenerator
 from .dead_code_eliminator import DeadCodeEliminator
 from .llvm_ir_generator import LLVMIRGenerator
 from .interleaver import Interleaver
+from .typedtree_printer import TypedtreePrinter
diff --git a/artiq/compiler/transforms/typedtree_printer.py b/artiq/compiler/transforms/typedtree_printer.py
new file mode 100644
index 000000000..e0c296f42
--- /dev/null
+++ b/artiq/compiler/transforms/typedtree_printer.py
@@ -0,0 +1,89 @@
+"""
+:class:`TypedtreePrinter` prints a human-readable representation of typedtrees.
+"""
+
+from pythonparser import algorithm, ast
+from .. import types, asttyped
+
+class TypedtreePrinter(algorithm.Visitor):
+    def __init__(self):
+        self.str = None
+        self.level = None
+        self.last_nl = None
+        self.type_printer = None
+
+    def print(self, node):
+        try:
+            self.str = ""
+            self.level = 0
+            self.last_nl = 0
+            self.type_printer = types.TypePrinter()
+            self.visit(node)
+            self._nl()
+            return self.str
+        finally:
+            self.str = None
+            self.level = None
+            self.last_nl = 0
+            self.type_printer = None
+
+    def _nl(self):
+        # self.str += "·"
+        if len(self.str) != self.last_nl:
+            self.str += "\n" + ("  " * self.level)
+            self.last_nl = len(self.str)
+
+    def _indent(self):
+        self.level += 1
+        self._nl()
+
+    def _dedent(self):
+        self._nl()
+        self.level -= 1
+        self.str = self.str[:-2]
+        self.last_nl -= 2
+
+    def visit(self, obj):
+        if isinstance(obj, ast.AST):
+            attrs = set(obj._fields) - {'ctx'}
+            if isinstance(obj, asttyped.commontyped):
+                attrs.update(set(obj._types))
+
+            for attr in set(attrs):
+                if not getattr(obj, attr):
+                    attrs.remove(attr) # omit falsey stuff
+
+            self.str += obj.__class__.__name__ + "("
+            if len(attrs) > 1:
+                self._indent()
+
+            for attr in attrs:
+                if len(attrs) > 1:
+                    self._nl()
+                self.str += attr + "="
+                self.visit(getattr(obj, attr))
+                if len(attrs) > 1:
+                    self._nl()
+
+            if len(attrs) > 1:
+                self._dedent()
+            self.str += ")"
+        elif isinstance(obj, types.Type):
+            self.str += self.type_printer.name(obj, max_depth=0)
+        elif isinstance(obj, list):
+            self.str += "["
+            if len(obj) > 1:
+                self._indent()
+
+            for elem in obj:
+                if len(obj) > 1:
+                    self._nl()
+                self.visit(elem)
+                if len(obj) > 1:
+                    self._nl()
+
+            if len(obj) > 1:
+                self._dedent()
+            self.str += "]"
+        else:
+            self.str += repr(obj)