1 files changed, 25 insertions, 9 deletions
diff --git a/railties/lib/commands/performance/profiler.rb b/railties/lib/commands/performance/profiler.rb
index 310d6764f1..464cea344d 100644
--- a/railties/lib/commands/performance/profiler.rb
+++ b/railties/lib/commands/performance/profiler.rb
@@ -1,5 +1,5 @@
 if ARGV.empty?
-  $stderr.puts "Usage: ./script/performance/profiler 'Person.expensive_method(10)' [times]"
+  $stderr.puts "Usage: ./script/performance/profiler 'Person.expensive_method(10)' [times] [flat|graph|graph_html]"
   exit(1)
 end
 
@@ -16,14 +16,30 @@ end
 
 # Use the ruby-prof extension if available.  Fall back to stdlib profiler.
 begin
-  require 'prof'
-  $stderr.puts 'Using the ruby-prof extension.'
-  Prof.clock_mode = Prof::GETTIMEOFDAY
-  Prof.start
-  profile_me
-  results = Prof.stop
-  require 'rubyprof_ext'
-  Prof.print_profile(results, $stderr)
+  begin
+    require "ruby-prof"
+    $stderr.puts 'Using the ruby-prof extension.'
+    RubyProf.clock_mode = RubyProf::WALL_TIME
+    RubyProf.start
+    profile_me
+    results = RubyProf.stop
+    if ARGV[2]
+      printer_class = RubyProf.const_get((ARGV[2] + "_printer").classify)
+    else
+      printer_class = RubyProf::FlatPrinter
+    end
+    printer = printer_class.new(results)
+    printer.print($stderr, 0)
+  rescue LoadError
+    require "prof"
+    $stderr.puts 'Using the old ruby-prof extension.'
+    Prof.clock_mode = Prof::GETTIMEOFDAY
+    Prof.start
+    profile_me
+    results = Prof.stop
+    require 'rubyprof_ext'
+    Prof.print_profile(results, $stderr)
+  end
 rescue LoadError
   require 'profiler'
   $stderr.puts 'Using the standard Ruby profiler.'