Commit a27a31c5 authored by Jan Hoenig's avatar Jan Hoenig
Browse files

Added a dot-printer. It is in pystencils/backend/dot.

parent e180d077
......@@ -36,11 +36,8 @@ class Node(object):
result.update(arg.atoms(argType))
return result
def parents(self):
return None
def children(self):
return None
yield None
class KernelFunction(Node):
......@@ -75,6 +72,7 @@ class KernelFunction(Node):
def __init__(self, body, functionName="kernel"):
super(KernelFunction, self).__init__()
self._body = body
body.parent = self
self._parameters = None
self._functionName = functionName
self._body.parent = self
......@@ -116,10 +114,14 @@ class KernelFunction(Node):
def children(self):
yield self.body
def __repr__(self):
def __str__(self):
self._updateParameters()
return '{0} {1}({2})\n{3}'.format(type(self).__name__, self.functionName, self.parameters,
textwrap.indent(repr(self.body), '\t'))
textwrap.indent(str(self.body), '\t'))
def __repr__(self):
self._updateParameters()
return '{0} {1}({2})'.format(type(self).__name__, self.functionName, self.parameters)
class Block(Node):
......@@ -179,10 +181,13 @@ class Block(Node):
return result - definedSymbols
def children(self):
yield self._nodes
return self._nodes
def __str__(self):
return ''.join('{!s}\n'.format(node) for node in self._nodes)
def __repr__(self):
return ''.join('{!r}\n'.format(node) for node in self._nodes)
return ''.join('{!r}'.format(node) for node in self._nodes)
class PragmaBlock(Block):
......@@ -196,6 +201,7 @@ class LoopOverCoordinate(Node):
def __init__(self, body, coordinateToLoopOver, start, stop, step=1):
self._body = body
body.parent = self
self._coordinateToLoopOver = coordinateToLoopOver
self._begin = start
self._end = stop
......@@ -278,11 +284,14 @@ class LoopOverCoordinate(Node):
return self._coordinateToLoopOver
def children(self):
return self.body
yield self.body
def __repr__(self):
def __str__(self):
return 'loop:{!s} in {!s}:{!s}:{!s}\n{!s}'.format(self.loopCounterName, self.start, self.stop, self.step,
textwrap.indent(repr(self.body), '\t'))
textwrap.indent(str(self.body), '\t'))
def __repr__(self):
return 'loop {!s} from {!s} to {!s} step{!s}'.format(self.loopCounterName, self.start, self.stop, self.step)
class SympyAssignment(Node):
......
from sympy.printing.printer import Printer
from graphviz import Digraph
class DotPrinter(Printer):
"""
A printer which converts ast to DOT (graph description language).
"""
def __init__(self, **kwargs):
super().__init__()
self.dot = Digraph(**kwargs)
def _print_KernelFunction(self, function):
self.dot.node(repr(function))
self._print(function.body)
def _print_LoopOverCoordinate(self, loop):
self.dot.node(repr(loop))
self._print(loop.body)
def _print_Block(self, block):
for node in block.children():
self._print(node)
parent = block.parent
for node in block.children():
self.dot.edge(repr(parent), repr(node))
parent = node
def _print_SympyAssignment(self, assignment):
self.dot.node(repr(assignment))
def doprint(self, expr):
self._print(expr)
return self.dot.source
def dotprint(ast, view=False, **kwargs):
"""
Returns a string which can be used to generate a DOT-graph
:param ast: The ast which should be generated
:param view: Boolen, if rendering of the image directly should occur.
:param kwargs: is directly passed to the DotPrinter class
:return: string in DOT format
"""
printer = DotPrinter(**kwargs)
dot = printer.doprint(ast)
if view:
printer.dot.render(view=view)
return dot
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment