dirtree in dirtree-0.6.0

- old
+ new

@@ -2,10 +2,11 @@
 require 'dirtree'
 require 'optparse'
 require 'json'
 require 'erb'
 require 'open-uri'
+require 'tempfile'
 
 templates_dir = File.join(File.dirname(__FILE__), '..', 'templates')
 templates = Dir.open(templates_dir).map do |file|
   file[0...file.index('.')]
 end.reject(&:empty?)
@@ -42,15 +43,23 @@
   ) do
     options[:local_dependencies] = true
   end
 
   opts.on(
+    '-s',
+    '--screenshot',
+    'Get an image screenshot of the directory tree'
+  ) do
+    options[:screenshot] = true
+  end
+
+  opts.on(
     '-oFile.html',
     '--output=File.html',
-    'Specify a path to write output, if
-       not specified output will be printed to STDOUT'
+    'Specify a path to write HTML output'
   ) do |value|
+    raise 'missing filename after ‘-o’' if value.empty?
     options[:output] = value
   end
 
   opts.on(
     '-tTemplateName',
@@ -80,10 +89,27 @@
 
     injection
   end
 end
 
-if options.key?(:output)
-  File.write(options[:output], result)
-else
-  puts result
+File.write(options[:output], result) if options.key?(:output)
+
+if options.key?(:screenshot)
+  unless options.key?(:output)
+    tmp = Tempfile.new
+    tmp.write(result)
+    tmp.close
+  end
+  file = options[:output] || tmp.path
+  %w[chromium chrome google-chrome google-chrome-stable].each do |chrome|
+    system(
+      chrome,
+      '--disable-gpu',
+      '--headless',
+      '--screenshot',
+      '--window-size=1500,1500',
+      "file://" + file
+    ) && break
+  end
 end
+
+puts result unless options.key?(:output) || options.key?(:screenshot)