bin/graphshaper in graphshaper-0.2.1 vs bin/graphshaper in graphshaper-0.2.2

@@ -1,40 +1,95 @@
 #!/usr/bin/env ruby
 
 require 'graphshaper'
+require 'optparse'
 
-if ARGV.length == 0 || ARGV[0] == "--help"
-  puts "  Usage: graphshaper SIZE"
-elsif ARGV[0].to_i < 21
-  puts "Please choose a size of at least 21"
-else
-  number_of_vertices = ARGV[0].to_i
-  inner_vertices = 20
+options = {}
+
+optparse = OptionParser.new do |opts|
+  opts.banner = "Usage: graphshaper [options] size"
   
-  vertex_output_file = File.new "vertices.csv", "w"
-  edge_output_file = File.new "edges.csv", "w"
-  logger = Graphshaper::LoggingAdapter.new vertex_output_file, edge_output_file
+  opts.on("-a", "--avocado", "Store the graph in a local AvocadoDB instance") do |avocado|
+    options[:avocado] = avocado
+  end
   
-  avocado_db = Graphshaper::AvocadoDbAdapter.new "vertices", "edges"
+  opts.on("-l", "--log", "Store the graph in two CSV files for nodes and edges") do |log|
+    options[:log] = log
+  end
   
-  start_time = Time.now
-  graph = Graphshaper::UndirectedGraph.without_orphans_with_order_of inner_vertices, adapters: [logger, avocado_db]
-  (number_of_vertices - inner_vertices).times do
-    graph.add_vertex do |preferential_attachment|
-      preferential_attachment > rand
-    end
+  opts.on("-d", "--dot", "Store the graph in the dot format") do |dot|
+    options[:dot] = dot
   end
-  end_time = Time.now
   
-  ellapsed_time = end_time - start_time
+  opts.on("-p", "--png", "Export the graph as a PNG (graphviz must be installed)") do |png|
+    options[:dot] = png
+    options[:png] = png
+  end
   
-  puts "#{graph.order} vertices (saved to vertices.csv)"
-  puts "#{graph.size} edges (saved to edges.csv)"
-  
-  if ellapsed_time < 2
-    puts "Generated and saved in about one second"
-  elsif ellapsed_time < 60
-    puts "Generated and saved in about #{ellapsed_time.round} seconds"
-  else
-    puts "Generated and saved in about #{ellapsed_time.round / 60} minutes and #{ellapsed_time.round % 60} seconds"
+  opts.on_tail("--version", "Show version") do
+    puts Graphshaper::VERSION
+    exit
   end
+  
+  opts.parse!
 end
+
+if ARGV.length == 0
+  puts "Please enter a size for your graph."
+  exit
+end
+
+number_of_vertices = ARGV.first.to_i
+inner_vertices = 20
+  
+adapters = []
+open_files = []
+  
+if options[:log]
+  vertex_output_file = File.new "vertices.csv", "w"
+  open_files << vertex_output_file
+  edge_output_file = File.new "edges.csv", "w"
+  open_files << edge_output_file
+  adapters << Graphshaper::LoggingAdapter.new(vertex_output_file, edge_output_file)
+end
+
+if options[:avocado]
+  adapters << Graphshaper::AvocadoDbAdapter.new("vertices", "edges")
+end
+
+if options[:dot]
+  dot_output_file = File.new "generated_graph.dot", "w"
+  open_files << dot_output_file
+  dot_adapter = Graphshaper::DotAdapter.new(dot_output_file)
+  adapters << dot_adapter
+end
+
+start_time = Time.now
+
+graph = Graphshaper::UndirectedGraph.without_orphans_with_order_of inner_vertices, adapters: adapters
+(number_of_vertices - inner_vertices).times do
+  graph.add_vertex { |preferential_attachment| preferential_attachment > rand }
+end
+
+dot_adapter.close if options[:dot]
+
+ellapsed_time = Time.now - start_time
+
+print "#{graph.order} vertices and #{graph.size} edges generated"
+
+print " and saved into AvocadoDB" if options[:avocado]
+print " and logged" if options[:log]
+print " and generated as a dot" if options[:dot]
+
+open_files.each { |file| file.close }
+
+if options[:png]
+  system('circo -Tpng generated_graph.dot -o generated_graph.png')
+end
+
+if ellapsed_time < 2
+  puts " in about one second"
+elsif ellapsed_time < 60
+  puts " in about #{ellapsed_time.round} seconds"
+else
+  puts " in about #{ellapsed_time.round / 60} minutes and #{ellapsed_time.round % 60} seconds"
+end
\ No newline at end of file