From 2d1bf75e8a2a92fa082f7103674a94c150405871 Mon Sep 17 00:00:00 2001
From: DannyB <dannyb@telesofia.com>
Date: Fri, 13 Jan 2017 10:59:59 +0200
Subject: [PATCH] fix graph size issues

---
 .gitignore                   |  1 +
 lib/decisiontree/id3_tree.rb | 12 ++++++++++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index d87d4be..7cc0635 100644
--- a/.gitignore
+++ b/.gitignore
@@ -15,3 +15,4 @@ spec/reports
 test/tmp
 test/version_tmp
 tmp
+*.png
\ No newline at end of file
diff --git a/lib/decisiontree/id3_tree.rb b/lib/decisiontree/id3_tree.rb
index b0702c8..c14f238 100755
--- a/lib/decisiontree/id3_tree.rb
+++ b/lib/decisiontree/id3_tree.rb
@@ -130,7 +130,13 @@ module DecisionTree
     def graph(filename, file_type = 'png')
       require 'graphr'
       dgp = DotGraphPrinter.new(build_tree)
+      dgp.size = ''
+      dgp.node_labeler = proc { |n| n.split("\n").first }
       dgp.write_to_file("#{filename}.#{file_type}", file_type)
+    rescue LoadError
+      STDERR.puts "Error: Cannot generate graph."
+      STDERR.puts "       The 'graphr' gem doesn't seem to be installed."
+      STDERR.puts "       Run 'gem install graphr' or add it to your Gemfile."
     end
 
     def ruleset
@@ -188,9 +194,11 @@ module DecisionTree
           child = attr[1][key]
           child_text = "#{child}\n(#{child.to_s.clone.object_id})"
         end
-        label_text = "#{key} ''"
+
         if type(attr[0].attribute) == :continuous
-          label_text.gsub!("''", attr[0].threshold.to_s)
+          label_text = "#{key} #{attr[0].threshold}"
+        else
+          label_text = key
         end
 
         [parent_text, child_text, label_text]