Allow filtering allocations by source file

[segiddins]

This is one of the most useful options of memory_profiler to my work on optimizing Bundler/RubyGems, since I'm usually focused on a single file/small set of files at a time, and seeing allocations from every file is less useful, especially if they far outnumber what I'm focused on. Additionally, early filtering of objects during the parse stage might help save memory in analysis

[byroot]

That makes sense to me.

I'm a bit busy at the moment, but I'll try to take a look at that in the near future, unless you wish to submit a PR?

[segiddins]

I'll see if I can get around to it, but no promises. This is what I'm running with now, it's just... very hacky.

--- tar:heap-profiler-0.7.0:lib/heap_profiler/analyzer.rb
+++ disk:heap-profiler-0.7.0:lib/heap_profiler/analyzer.rb
@@ -32,12 +32,16 @@ module HeapProfiler
           klass = case grouping
           when "file"
             FileGroupDimension
+          when "directory"
+            DirectoryGroupDimension
           when "location"
             LocationGroupDimension
           when "gem"
             GemGroupDimension
           when "class"
             ClassGroupDimension
+          when "class_location"
+            ClassLocationGroupDimension
           else
             raise "Unknown grouping key: #{grouping.inspect}"
           end
@@ -92,6 +96,23 @@ module HeapProfiler
       end
     end
 
+    class DirectoryGroupDimension < GroupedDimension
+      def initialize(...)
+        super
+        @grouping = Hash.new { |h, k| h[k] = File.dirname(k) }
+      end
+
+      def process(_index, object)
+        file = object[:file]
+
+        if file
+          grouping = @grouping[file]
+          @objects[grouping] += 1
+          @memory[grouping] += object[:memsize]
+        end
+      end
+    end
+
     class GemGroupDimension < GroupedDimension
       def process(index, object)
         if (group = index.guess_gem(object))
@@ -110,6 +131,19 @@ module HeapProfiler
       end
     end
 
+    class ClassLocationGroupDimension < GroupedDimension
+      def process(index, object)
+        file = object[:file]
+        line = object[:line]
+
+        if file && line && (klass = index.guess_class(object))
+          group = "#{file}:#{line} - #{klass}"
+          @objects[group] += 1
+          @memory[group] += object[:memsize]
+        end
+      end
+    end
+
     class StringDimension
       class StringLocation
         attr_reader :location, :count, :memsize
@@ -223,7 +257,9 @@ module HeapProfiler
       end
 
       processors = dimensions.values
+      allow_files = ENV["ALLOW_FILES"] && Regexp.new(ENV["ALLOW_FILES"])
       @heap.each_object do |object|
+        next if allow_files && !allow_files.match?(object[:file])
         processors.each { |p| p.process(@index, object) }
       end
       dimensions
--- tar:heap-profiler-0.7.0:lib/heap_profiler/cli.rb
+++ disk:heap-profiler-0.7.0:lib/heap_profiler/cli.rb
@@ -44,7 +44,13 @@ module HeapProfiler
       else
         HeapResults.new(path)
       end
-      results.pretty_print(scale_bytes: true, normalize_paths: true)
+      if @output_file && @output_file != "-"
+        File.open(@output_file, "w") do |output|
+          results.pretty_print(output, scale_bytes: true, normalize_paths: true)
+        end
+      else
+        results.pretty_print(scale_bytes: true, normalize_paths: true)
+      end
     end
 
     def clean_dump(path)
@@ -118,6 +124,10 @@ module HeapProfiler
 
         opts.on('-r', '--retained-only', 'Only compute report for memory retentions.') do
           @retained_only = true
+        end
+
+        opts.on('-o', '--output=FILE', 'Write report to FILE instead of STDOUT.') do |arg|
+          @output_file = arg
         end
 
         HeapProfiler::AbstractResults.top_entries_count = 50
--- tar:heap-profiler-0.7.0:lib/heap_profiler/index.rb
+++ disk:heap-profiler-0.7.0:lib/heap_profiler/index.rb
@@ -75,6 +75,8 @@ module HeapProfiler
       @gems[path] ||=
         if %r{(/gems/.*)*/gems/(?<gemname>[^/]+)} =~ path
           gemname
+        elsif %r{/lib/bundler[\./]}.match?(path)
+          "bundler"
         elsif %r{/rubygems[\./]}.match?(path)
           "rubygems"
         elsif %r{ruby/2\.[^/]+/(?<stdlib>[^/\.]+)} =~ path
--- tar:heap-profiler-0.7.0:lib/heap_profiler/results.rb
+++ disk:heap-profiler-0.7.0:lib/heap_profiler/results.rb
@@ -16,7 +16,7 @@ module HeapProfiler
 
     METRICS = ["memory", "objects", "strings", "shape_edges"].freeze
     GROUPED_METRICS = ["memory", "objects"]
-    GROUPINGS = ["gem", "file", "location", "class"].freeze
+    GROUPINGS = ["gem", "directory", "file", "location", "class", "class_location"].freeze
 
     attr_reader :types, :dimensions