diff --git a/backends/dot.py b/backends/dot.py index d477da0a4d0cc5ec5225e929bc17f03577c68b48..52184a76659708657f4341ab065dc2dda075af14 100644 --- a/backends/dot.py +++ b/backends/dot.py @@ -8,10 +8,9 @@ class DotPrinter(Printer): """ A printer which converts ast to DOT (graph description language). """ - def __init__(self, node_to_str_function, full, **kwargs): + def __init__(self, node_to_str_function, **kwargs): super(DotPrinter, self).__init__() self._node_to_str_function = node_to_str_function - self.full = full self.dot = Digraph(**kwargs) self.dot.quote_edge = lang.quote @@ -36,11 +35,6 @@ class DotPrinter(Printer): def _print_SympyAssignment(self, assignment): self.dot.node(str(id(assignment)), style='filled', fillcolor='#56db7f', label=self._node_to_str_function(assignment)) - if self.full: - for node in assignment.args: - self._print(node) - for node in assignment.args: - self.dot.edge(str(id(assignment)), str(id(node))) def _print_Conditional(self, expr): self.dot.node(str(id(expr)), style='filled', fillcolor='#56bd7f', label=self._node_to_str_function(expr)) @@ -50,16 +44,6 @@ class DotPrinter(Printer): self._print(expr.false_block) self.dot.edge(str(id(expr)), str(id(expr.false_block))) - def empty_printer(self, expr): - if self.full: - self.dot.node(str(id(expr)), label=self._node_to_str_function(expr)) - for node in expr.args: - self._print(node) - for node in expr.args: - self.dot.edge(str(id(expr)), str(id(node))) - else: - raise NotImplementedError('DotPrinter cannot print', type(expr), expr) - def doprint(self, expr): self._print(expr) return self.dot.source @@ -83,23 +67,19 @@ def __shortened(node): raise NotImplementedError("Cannot handle node type %s" % (type(node),)) -def print_dot(node, view=False, short=False, full=False, **kwargs): +def print_dot(node, view=False, short=False, **kwargs): """ Returns a string which can be used to generate a DOT-graph :param node: The ast which should be generated :param view: Boolean, if rendering of the image directly should occur. :param short: Uses the __shortened output - :param full: Prints the whole tree with type information :param kwargs: is directly passed to the DotPrinter class: http://graphviz.readthedocs.io/en/latest/api.html#digraph :return: string in DOT format """ node_to_str_function = repr if short: node_to_str_function = __shortened - elif full: - def node_to_str_function(expr): - return repr(type(expr)) + repr(expr) - printer = DotPrinter(node_to_str_function, full, **kwargs) + printer = DotPrinter(node_to_str_function, **kwargs) dot = printer.doprint(node) if view: return graphviz.Source(dot) diff --git a/display_utils.py b/display_utils.py index 2d0fa87f7cb872c9ae0bdfaa404e2bb27764fe58..3607e45b99bd1f6cc81415a05ad300ef344672fc 100644 --- a/display_utils.py +++ b/display_utils.py @@ -3,7 +3,7 @@ from typing import Any, Dict, Optional from pystencils.astnodes import KernelFunction -def to_dot(expr: sp.Expr, graph_style: Optional[Dict[str, Any]] = None): +def to_dot(expr: sp.Expr, graph_style: Optional[Dict[str, Any]] = None, short=True): """Show a sympy or pystencils AST as dot graph""" from pystencils.astnodes import Node import graphviz @@ -11,7 +11,7 @@ def to_dot(expr: sp.Expr, graph_style: Optional[Dict[str, Any]] = None): if isinstance(expr, Node): from pystencils.backends.dot import print_dot - return graphviz.Source(print_dot(expr, short=True, graph_attr=graph_style)) + return graphviz.Source(print_dot(expr, short=short, graph_attr=graph_style)) else: from sympy.printing.dot import dotprint return graphviz.Source(dotprint(expr, graph_attr=graph_style))