Custom printer support

bin/ruby-prof

@@ -6,6 +6,7 @@ require 'ruby-prof'
 # Now setup option parser
 require 'ostruct'
 require 'optparse'
+require 'facets/string/modulize'
 
 module RubyProf
   # == Synopsis
@@ -53,6 +54,8 @@ module RubyProf
     attr_accessor :options
     attr_reader :profile
 
+    BUILTIN_PRINTERS = [:flat, :flat_with_line_numbers, :graph, :graph_html, :call_tree, :call_stack, :dot, :multi].freeze
+
     def initialize
       setup_options
       parse_args
@@ -118,35 +121,34 @@ module RubyProf
         opts.separator ""
         opts.separator "Options:"
 
-        opts.on('-p printer', '--printer=printer', [:flat, :flat_with_line_numbers, :graph, :graph_html, :call_tree, :call_stack, :dot, :multi],
+        opts.on('-p printer', '--printer=printer',
                 'Select a printer:',
                 '  flat - Prints a flat profile as text (default).',
                 '  graph - Prints a graph profile as text.',
                 '  graph_html - Prints a graph profile as html.',
                 '  call_tree - format for KCacheGrind',
                 '  call_stack - prints a HTML visualization of the call tree',
                 '  dot - Prints a graph profile as a dot file',
-                '  multi - Creates several reports in output directory'
+                '  multi - Creates several reports in output directory',
+                '  <class name> - Use a custom class'
         ) do |printer|
 
-          case printer
-          when :flat
-            options.printer = RubyProf::FlatPrinter
-          when :graph
-            options.printer = RubyProf::GraphPrinter
-          when :graph_html
-            options.printer = RubyProf::GraphHtmlPrinter
-          when :call_tree
-            options.printer = RubyProf::CallTreePrinter
-          when :call_stack
-            options.printer = RubyProf::CallStackPrinter
-          when :dot
-            options.printer = RubyProf::DotPrinter
-          when :multi
-            options.printer = RubyProf::MultiPrinter
+
+          if BUILTIN_PRINTERS.include?(printer.to_sym)
+            options.printer_class_name = "RubyProf::#{printer}Printer".modulize
+          elsif printer.modulize == printer # if this is already in a modularized, it's a class
+            options.printer_class_name = printer
+          else
+            puts "Invalid printer: #{printer}"
+            puts opts
+            exit(-1)
           end
         end
 
+        opts.on('--printer-require=require', 'A ruby class to load for loading the printer') do |require|
+          options.printer_require = require
+        end
+
         opts.on('-m min_percent', '--min_percent=min_percent', Float,
                 'The minimum percent a method must take before ',
                 '  being included in output reports.',
@@ -242,6 +244,31 @@ module RubyProf
       end
     end
 
+    def load_printer
+      if options.printer_require
+        begin
+          require options.printer_require
+        rescue LoadError
+          puts "Couldn't require #{options.printer_require} Double check that it is correct, and any gems needed are installed"
+          puts
+          puts option_parser
+          exit(-1)
+        end
+      end
+
+      # class name given? constantize it now
+      if options.printer_class_name
+        begin
+          options.printer = constantize(options.printer_class_name)
+        rescue NameError
+          puts "Couldn't find #{options.printer_class_name} to use as a printer. Double check that it is correct, and that you have given a --printer-require if necessary."
+          puts
+          puts option_parser
+          exit(-1)
+        end
+      end
+    end
+
     def parse_args
       # Make sure the user specified at least one file
       if ARGV.length < 1 and not options.exec
@@ -253,6 +280,9 @@ module RubyProf
 
       self.option_parser.parse! ARGV
 
+      # make sure it's loaded from options in time to check for directory
+      load_printer
+
       if options.printer.needs_dir?
         options.file ||= "."
         options.old_wd ||= Dir.pwd

ruby-prof.gemspec

@@ -59,6 +59,7 @@ EOF
   spec.required_ruby_version = '>= 2.7.0'
   spec.date = Time.now.strftime('%Y-%m-%d')
   spec.homepage = 'https://github.com/ruby-prof/ruby-prof'
+  spec.add_dependency("facets")
   spec.add_development_dependency('minitest')
   spec.add_development_dependency('rake-compiler')
 end