allow :dirs option for .enumerate

Allows custom rake tasks to be defined using:

  SourceAnnotationExtractor.enumerate 'TODO|FIXME', dirs: %w(app lib), tag: true

3 files changed, 60 insertions, 6 deletions

diff --git a/railties/CHANGELOG.md b/railties/CHANGELOG.md
index 4fd1a2369f..4f18f4cdd2 100644
--- a/railties/CHANGELOG.md
+++ b/railties/CHANGELOG.md
@@ -1,5 +1,11 @@
 ## Rails 4.0.0 (unreleased) ##
 
+*   Allow a `:dirs` key in the `SourceAnnotationExtractor.enumerate` options
+    to explicitly set the directories to be traversed so it's easier to define
+    custom rake tasks.
+
+    *Brian D. Burns*
+
 *   Deprecate `Rails::Generators::ActiveModel#update_attributes` in favor of `#update`.
 
     ORMs that implement `Generators::ActiveModel#update_attributes` should change
diff --git a/railties/lib/rails/source_annotation_extractor.rb b/railties/lib/rails/source_annotation_extractor.rb
index 0a352f1a18..2cbb0a435c 100644
--- a/railties/lib/rails/source_annotation_extractor.rb
+++ b/railties/lib/rails/source_annotation_extractor.rb
@@ -32,15 +32,24 @@ class SourceAnnotationExtractor
   end
 
   # Prints all annotations with tag +tag+ under the root directories +app+,
-  # +config+, +lib+, and +test+ (recursively). Filenames with extension
-  # +.builder+, +.rb+, +.erb+, +.haml+, +.slim+, +.css+, +.scss+, +.js+,
-  # +.coffee+, and +.rake+ are taken into account. The +options+ hash is
-  # passed to each annotation's +to_s+.
+  # +config+, +db+, +lib+, and +test+ (recursively).
+  #
+  # Additional directories may be added using a comma-delimited list set using
+  # <tt>ENV['SOURCE_ANNOTATION_DIRECTORIES']</tt>.
+  #
+  # Directories may also be explicitly set using the <tt>:dirs</tt> key in +options+.
+  #
+  #   SourceAnnotationExtractor.enumerate 'TODO|FIXME', dirs: %w(app lib), tag: true
+  #
+  # If +options+ has a <tt>:tag</tt> flag, it will be passed to each annotation's +to_s+.
+  #
+  # See <tt>#find_in</tt> for a list of file extensions that will be taken into account.
   #
   # This class method is the single entry point for the rake tasks.
   def self.enumerate(tag, options={})
     extractor = new(tag)
-    extractor.display(extractor.find, options)
+    dirs = options.delete(:dirs) || Annotation.directories
+    extractor.display(extractor.find(dirs), options)
   end
 
   attr_reader :tag
@@ -51,7 +60,7 @@ class SourceAnnotationExtractor
 
   # Returns a hash that maps filenames under +dirs+ (recursively) to arrays
   # with their annotations.
-  def find(dirs = Annotation.directories)
+  def find(dirs)
     dirs.inject({}) { |h, dir| h.update(find_in(dir)) }
   end
 
diff --git a/railties/test/application/rake/notes_test.rb b/railties/test/application/rake/notes_test.rb
index 5893d58925..3508f4225a 100644
--- a/railties/test/application/rake/notes_test.rb
+++ b/railties/test/application/rake/notes_test.rb
@@ -130,6 +130,45 @@ module ApplicationTests
         end
       end
 
+      test 'custom rake task finds specific notes in specific directories' do
+        app_file "app/controllers/some_controller.rb", "# TODO: note in app directory"
+        app_file "lib/some_file.rb", "# OPTIMIZE: note in lib directory\n" << "# FIXME: note in lib directory"
+        app_file "test/some_test.rb", 1000.times.map { "" }.join("\n") << "# TODO: note in test directory"
+
+        app_file "lib/tasks/notes_custom.rake", <<-EOS
+          require 'rails/source_annotation_extractor'
+          task :notes_custom do
+            tags = 'TODO|FIXME'
+            opts = { dirs: %w(lib test), tag: true }
+            SourceAnnotationExtractor.enumerate(tags, opts)
+          end
+        EOS
+
+        boot_rails
+
+        require 'rake'
+        require 'rdoc/task'
+        require 'rake/testtask'
+
+        Rails.application.load_tasks
+
+        Dir.chdir(app_path) do
+          output = `bundle exec rake notes_custom`
+          lines = output.scan(/\[([0-9\s]+)\]/).flatten
+
+          assert_match(/\[FIXME\] note in lib directory/, output)
+          assert_match(/\[TODO\] note in test directory/, output)
+          assert_no_match(/OPTIMIZE/, output)
+          assert_no_match(/note in app directory/, output)
+
+          assert_equal 2, lines.size
+
+          lines.each do |line_number|
+            assert_equal 4, line_number.size
+          end
+        end
+      end
+
       private
       def boot_rails
         super