lib/hotch.rb in hotch-0.0.1 vs lib/hotch.rb in hotch-0.0.2

@@ -1,32 +1,92 @@
 require 'stackprof'
 require 'tmpdir'
 
-HOTCH_VIEWER = ENV['HOTCH_VIEWER']
+class Hotch
+  attr_reader :name
 
-NAME = $0.gsub(/\W/, '_')
+  def initialize(name)
+    @name = name
+    @reports = []
+  end
 
-StackProf.start
+  def start(*args)
+    StackProf.start(*args) unless StackProf.running?
+  end
 
-def system!(*args)
-  system(*args) or raise "system call failed: #{args.join(' ')}"
-end
+  def stop
+    if StackProf.running?
+      StackProf.stop
+      @reports << StackProf::Report.new(results)
+    end
+  end
 
-at_exit do
-  StackProf.stop
+  def run(*args)
+    start(*args)
+    yield
+  ensure
+    stop
+  end
 
-  dir = Dir.mktmpdir("hotch.#{NAME}")
-  dump = File.open(File.join(dir, "profile.dump"), "wb")
-  svg = File.join(File.join(dir, "profile.svg"))
+  def results
+    StackProf.results
+  end
 
-  StackProf::Report.new(StackProf.results).print_graphviz(nil, dump)
+  def report
+    report = @reports.inject(:+) or return
 
-  dump.close
+    dir = Dir.mktmpdir("hotch.#{name}.")
+    dump = File.open(File.join(dir, "profile.dump"), "wb")
+    svg = File.join(File.join(dir, "profile.svg"))
 
-  system! "dot", "-Tsvg", "-o", svg, dump.path
 
-  puts "Profile SVG: #{svg}"
+    report.print_graphviz(nil, dump)
 
-  if viewer = HOTCH_VIEWER
-    system! viewer, svg
+    dump.close
+
+    system! "dot", "-Tsvg", "-o", svg, dump.path
+
+    yield svg
+  end
+
+  def system!(*args)
+    Kernel.system(*args) or raise "system call failed: #{args.join(' ')}"
+  end
+
+  def report_at_exit(viewer=ENV['HOTCH_VIEWER'])
+    return if defined? @at_exit_installed
+
+    at_exit do
+      stop
+
+      report do |svg|
+        puts "Profile SVG: #{svg}"
+
+        if viewer
+          system viewer, svg
+        end
+      end
+    end
+
+    @at_exit_installed = true
+  end
+
+  private
+
+  def name
+    @name.gsub(/\W+/, '_')
+  end
+end
+
+def Hotch(aggregate: true)
+  hotch = if aggregate
+    $hotch ||= Hotch.new($0)
+  else
+    caller = Kernel.caller_locations(1).first
+    Hotch.new("#$0:#{caller.path}:#{caller.lineno}")
+  end
+
+  hotch.report_at_exit
+  hotch.run do
+    yield
   end
 end