Add Public Api for register new extensions for Rake Notes

2 files changed, 29 insertions, 15 deletions

diff --git a/railties/lib/rails/source_annotation_extractor.rb b/railties/lib/rails/source_annotation_extractor.rb
index 83e28090f8..5863ad1ed7 100644
--- a/railties/lib/rails/source_annotation_extractor.rb
+++ b/railties/lib/rails/source_annotation_extractor.rb
@@ -18,6 +18,22 @@ class SourceAnnotationExtractor
       @@directories ||= %w(app config db lib test) + (ENV['SOURCE_ANNOTATION_DIRECTORIES'] || '').split(',')
     end
 
+    def self.extensions
+      @@extensions ||= {}
+    end
+
+    # Registers new Annotations File Extensions
+    #   SourceAnnotationExtractor::Annotation.register_extensions("css", "scss", "sass", "less", "js") { |tag| /\/\/\s*(#{tag}):?\s*(.*)$/ }
+    def self.register_extensions(*extensions, &block)
+      self.extensions[/\.(#{extensions.join("|")})$/] = block
+    end
+
+    register_extensions("builder", "rb", "coffe", "rake") { |tag| /#\s*(#{tag}):?\s*(.*)$/ }
+    register_extensions("css", "scss", "sass", "less", "js") { |tag| /\/\/\s*(#{tag}):?\s*(.*)$/ }
+    register_extensions("erb") { |tag| /<%\s*#\s*(#{tag}):?\s*(.*?)\s*%>/ }
+    register_extensions("haml") { |tag| /-\s*#\s*(#{tag}):?\s*(.*)$/ }
+    register_extensions("slim") { |tag| /\/\s*\s*(#{tag}):?\s*(.*)$/ }
+
     # Returns a representation of the annotation that looks like this:
     #
     #   [126] [TODO] This algorithm is simple and clearly correct, make it faster.
@@ -78,21 +94,14 @@ class SourceAnnotationExtractor
       if File.directory?(item)
         results.update(find_in(item))
       else
-        pattern =
-            case item
-            when /\.(builder|rb|coffee|rake)$/
-              /#\s*(#{tag}):?\s*(.*)$/
-            when /\.(css|scss|sass|less|js)$/
-              /\/\/\s*(#{tag}):?\s*(.*)$/
-            when /\.erb$/
-              /<%\s*#\s*(#{tag}):?\s*(.*?)\s*%>/
-            when /\.haml$/
-              /-\s*#\s*(#{tag}):?\s*(.*)$/
-            when /\.slim$/
-              /\/\s*\s*(#{tag}):?\s*(.*)$/
-            else nil
-            end
-        results.update(extract_annotations_from(item, pattern)) if pattern
+        extension = Annotation.extensions.detect do |regexp, _block|
+          regexp.match(item)
+        end
+
+        if extension
+          pattern = extension.last.call(tag)
+          results.update(extract_annotations_from(item, pattern)) if pattern
+        end
       end
     end
 
diff --git a/railties/test/application/rake/notes_test.rb b/railties/test/application/rake/notes_test.rb
index 05f6338b68..62ab28082d 100644
--- a/railties/test/application/rake/notes_test.rb
+++ b/railties/test/application/rake/notes_test.rb
@@ -175,6 +175,11 @@ module ApplicationTests
         end
       end
 
+      test 'register a new extension' do
+        SourceAnnotationExtractor::Annotation.register_extensions(".test1", ".test2") { |tag| /#{tag}/ }
+        assert SourceAnnotationExtractor::Annotation.extensions[/(\.test1|\.test2)/]
+      end
+
       private
       def boot_rails
         super