From a01de190d3bff78de82fc4d86f8b29095f2ae03b Mon Sep 17 00:00:00 2001
From: Omega16 <22673084+omega16@users.noreply.github.com>
Date: Mon, 24 Aug 2020 15:35:29 -0500
Subject: [PATCH] Added pydot__tree_to_dot and pydot__tree_to_graph, changed
 pydot__tree_to_png

---
 lark/tree.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/lark/tree.py b/lark/tree.py
index b9dddf4..45346c6 100644
--- a/lark/tree.py
+++ b/lark/tree.py
@@ -160,6 +160,15 @@ class SlottedTree(Tree):
 
 
 def pydot__tree_to_png(tree, filename, rankdir="LR", **kwargs):
+    graph = pydot__tree_to_graph(tree, rankdir, **kwargs)
+    graph.write_png(filename)
+
+
+def pydot__tree_to_dot(tree, filename, rankdir="LR", **kwargs):
+    graph = pydot__tree_to_graph(tree, rankdir, **kwargs)
+    graph.write(filename)
+
+def pydot__tree_to_graph(tree, rankdir="LR", **kwargs):
     """Creates a colorful image that represents the tree (data+children, without meta)
 
     Possible values for `rankdir` are "TB", "LR", "BT", "RL", corresponding to
@@ -197,5 +206,5 @@ def pydot__tree_to_png(tree, filename, rankdir="LR", **kwargs):
         return node
 
     _to_pydot(tree)
-    graph.write_png(filename)
-
+    return graph
+