MAINT make export_graphviz more exception-safe

Author: larsmans

sklearn/tree/export.py

@@ -115,20 +115,22 @@ def recurse(tree, node_id, criterion, parent=None, depth=0):
                 out_file.write('%d -> %d ;\n' % (parent, node_id))
 
     own_file = False
-    if isinstance(out_file, six.string_types):
-        if six.PY3:
-            out_file = open(out_file, "w", encoding="utf-8")
-        else:
-            out_file = open(out_file, "wb")
-        own_file = True
+    try:
+        if isinstance(out_file, six.string_types):
+            if six.PY3:
+                out_file = open(out_file, "w", encoding="utf-8")
+            else:
+                out_file = open(out_file, "wb")
+            own_file = True
 
-    out_file.write("digraph Tree {\n")
+        out_file.write("digraph Tree {\n")
 
-    if isinstance(decision_tree, _tree.Tree):
-        recurse(decision_tree, 0, criterion="impurity")
-    else:
-        recurse(decision_tree.tree_, 0, criterion=decision_tree.criterion)
-    out_file.write("}")
+        if isinstance(decision_tree, _tree.Tree):
+            recurse(decision_tree, 0, criterion="impurity")
+        else:
+            recurse(decision_tree.tree_, 0, criterion=decision_tree.criterion)
+        out_file.write("}")
 
-    if own_file:
-        out_file.close()
+    finally:
+        if own_file:
+            out_file.close()