Support annotating model fixture files (#110)

This PR fixes adding annotation to YML files, like fixtures. The Ruby
parser introduced in #72 and #82 did not support YML files.

lib/annotate_rb/model_annotator.rb

@@ -27,5 +27,6 @@ module ModelAnnotator
     autoload :FileParser, "annotate_rb/model_annotator/file_parser"
     autoload :ZeitwerkClassGetter, "annotate_rb/model_annotator/zeitwerk_class_getter"
     autoload :CheckConstraintAnnotation, "annotate_rb/model_annotator/check_constraint_annotation"
+    autoload :FileToParserMapper, "annotate_rb/model_annotator/file_to_parser_mapper"
   end
 end

lib/annotate_rb/model_annotator/annotated_file/generator.rb

@@ -5,7 +5,7 @@ module ModelAnnotator
     module AnnotatedFile
       # Generates the file with fresh annotations
       class Generator
-        def initialize(file_content, new_annotations, annotation_position, options)
+        def initialize(file_content, new_annotations, annotation_position, parser_klass, parsed_file, options)
           @annotation_position = annotation_position
           @options = options
 
@@ -14,7 +14,8 @@ def initialize(file_content, new_annotations, annotation_position, options)
           @new_annotations = new_annotations
           @file_content = file_content
 
-          @parsed_file = FileParser::ParsedFile.new(@file_content, @new_annotations, options).parse
+          @parser = parser_klass
+          @parsed_file = parsed_file
         end
 
         # @return [String] Returns the annotated file content to be written back to a file
@@ -32,7 +33,7 @@ def generate
           end
 
           # We need to get class start and class end depending on the position
-          parsed = FileParser::CustomParser.new(content_without_annotations, "", 0).tap(&:parse)
+          parsed = @parser.parse(content_without_annotations)
 
           _content = if %w[after bottom].include?(annotation_write_position)
             content_annotated_after(parsed, content_without_annotations)

lib/annotate_rb/model_annotator/annotated_file/updater.rb

@@ -5,13 +5,13 @@ module ModelAnnotator
     module AnnotatedFile
       # Updates existing annotations
       class Updater
-        def initialize(file_content, new_annotations, _annotation_position, options)
+        def initialize(file_content, new_annotations, _annotation_position, parsed_file, options)
           @options = options
 
           @new_annotations = new_annotations
           @file_content = file_content
 
-          @parsed_file = FileParser::ParsedFile.new(@file_content, @new_annotations, options).parse
+          @parsed_file = parsed_file
         end
 
         # @return [String] Returns the annotated file content to be written back to a file

lib/annotate_rb/model_annotator/file_parser.rb

@@ -7,6 +7,7 @@ module FileParser
       autoload :CustomParser, "annotate_rb/model_annotator/file_parser/custom_parser"
       autoload :ParsedFile, "annotate_rb/model_annotator/file_parser/parsed_file"
       autoload :ParsedFileResult, "annotate_rb/model_annotator/file_parser/parsed_file_result"
+      autoload :YmlParser, "annotate_rb/model_annotator/file_parser/yml_parser"
     end
   end
 end

lib/annotate_rb/model_annotator/file_parser/annotation_finder.rb

@@ -19,19 +19,17 @@ class NoAnnotationFound < StandardError; end
 
         attr_reader :parser
 
-        def initialize(content, wrapper_open, wrapper_close)
+        def initialize(content, wrapper_open, wrapper_close, parser)
           @content = content
           @wrapper_open = wrapper_open
           @wrapper_close = wrapper_close
           @annotation_start = nil
           @annotation_end = nil
-          @parser = nil
+          @parser = parser
         end
 
         # Find the annotation's line start and line end
         def run
-          # CustomParser returns line numbers as 0-indexed
-          @parser = FileParser::CustomParser.new(@content, "", 0).tap(&:parse)
           comments = @parser.comments
 
           start = comments.find_index { |comment, _| comment.include?(COMPAT_PREFIX) || comment.include?(COMPAT_PREFIX_MD) }

lib/annotate_rb/model_annotator/file_parser/parsed_file.rb

@@ -6,15 +6,17 @@ module FileParser
       class ParsedFile
         SKIP_ANNOTATION_STRING = "# -*- SkipSchemaAnnotations"
 
-        def initialize(file_content, new_annotations, options)
+        def initialize(file_content, new_annotations, parser_klass, options)
           @file_content = file_content
           @file_lines = @file_content.lines
           @new_annotations = new_annotations
+          @parser_klass = parser_klass
           @options = options
         end
 
         def parse
-          @finder = AnnotationFinder.new(@file_content, @options[:wrapper_open], @options[:wrapper_close])
+          @file_parser = @parser_klass.parse(@file_content)
+          @finder = AnnotationFinder.new(@file_content, @options[:wrapper_open], @options[:wrapper_close], @file_parser)
           has_annotations = false
 
           begin
@@ -30,7 +32,6 @@ def parse
           end
 
           @diff = AnnotationDiffGenerator.new(annotations, @new_annotations).generate
-          @file_parser = @finder.parser
 
           has_skip_string = @file_parser.comments.any? { |comment, _lineno| comment.include?(SKIP_ANNOTATION_STRING) }
           annotations_changed = @diff.changed?
@@ -85,7 +86,9 @@ def parse
             annotations_with_whitespace: annotations_with_whitespace,
             has_leading_whitespace: has_leading_whitespace,
             has_trailing_whitespace: has_trailing_whitespace,
-            annotation_position: annotation_position
+            annotation_position: annotation_position,
+            starts: @file_parser.starts,
+            ends: @file_parser.ends
           )
         end
       end

lib/annotate_rb/model_annotator/file_parser/parsed_file_result.rb

@@ -12,7 +12,9 @@ def initialize(
           annotations_with_whitespace:,
           has_leading_whitespace:,
           has_trailing_whitespace:,
-          annotation_position:
+          annotation_position:,
+          starts:,
+          ends:
         )
           @has_annotations = has_annotations
           @has_skip_string = has_skip_string
@@ -22,9 +24,11 @@ def initialize(
           @has_leading_whitespace = has_leading_whitespace
           @has_trailing_whitespace = has_trailing_whitespace
           @annotation_position = annotation_position
+          @starts = starts
+          @ends = ends
         end
 
-        attr_reader :annotations, :annotation_position
+        attr_reader :annotations, :annotation_position, :starts, :ends
 
         # Returns annotations with new line before and after if they exist
         attr_reader :annotations_with_whitespace

lib/annotate_rb/model_annotator/file_parser/yml_parser.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+require "psych"
+
+module AnnotateRb
+  module ModelAnnotator
+    module FileParser
+      class YmlParser
+        class << self
+          def parse(string)
+            _parser = new(string).tap(&:parse)
+          end
+        end
+
+        attr_reader :comments, :starts, :ends
+
+        def initialize(input)
+          @input = input
+          @comments = []
+          @starts = []
+          @ends = []
+        end
+
+        def parse
+          parse_comments
+          parse_yml
+        end
+
+        private
+
+        def parse_comments
+          # Adds 0-indexed line numbers
+          @input.split($/).each_with_index do |line, line_no|
+            if line.strip.starts_with?("#")
+              @comments << [line, line_no]
+            end
+          end
+        end
+
+        def parse_yml
+          # https://docs.ruby-lang.org/en/master/Psych.html#module-Psych-label-Reading+to+Psych-3A-3ANodes-3A-3AStream+structure
+          parser = Psych.parser
+          parser.parse(@input)
+
+          stream = parser.handler.root
+
+          if stream.children.any?
+            doc = stream.children.first
+            @starts << [nil, doc.start_line]
+            @ends << [nil, doc.end_line]
+          else
+            # When parsing a yml file, streamer returns an instance of `Psych::Nodes::Stream` which is a subclass of
+            #   `Psych::Nodes::Node`. It along with children nodes, implement #start_line and #end_line.
+            #
+            # When parsing input that is only comments, the parser counts #start_line as the start of the file being
+            #   line 0.
+            #
+            # What we really want is where the "start" of the yml file would happen, which would be after comments.
+            # This is stream#end_line.
+            @starts << [nil, stream.end_line]
+            @ends << [nil, stream.end_line]
+          end
+        end
+      end
+    end
+  end
+end

lib/annotate_rb/model_annotator/file_to_parser_mapper.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module AnnotateRb
+  module ModelAnnotator
+    class FileToParserMapper
+      class UnsupportedFileTypeError < StandardError; end
+
+      MAP = {
+        ".rb" => FileParser::CustomParser,
+        ".yml" => FileParser::YmlParser
+      }.freeze
+
+      class << self
+        def map(file_name)
+          extension = File.extname(file_name).downcase
+          parser = MAP[extension]
+
+          raise UnsupportedFileTypeError, "File '#{file_name}' does not have a supported file type." if parser.nil?
+
+          parser
+        end
+      end
+    end
+  end
+end

lib/annotate_rb/model_annotator/single_file_annotation_remover.rb

@@ -12,8 +12,10 @@ def call(file_name, options = Options.from({}))
           return false unless File.exist?(file_name)
           old_content = File.read(file_name)
 
+          parser_klass = FileToParserMapper.map(file_name)
+
           begin
-            parsed_file = FileParser::ParsedFile.new(old_content, "", options).parse
+            parsed_file = FileParser::ParsedFile.new(old_content, "", parser_klass, options).parse
           rescue FileParser::AnnotationFinder::MalformedAnnotation => e
             warn "Unable to process #{file_name}: #{e.message}"
             warn "\t" + e.backtrace.join("\n\t") if @options[:trace]

lib/annotate_rb/model_annotator/single_file_annotator.rb

@@ -25,8 +25,10 @@ def call(file_name, annotation, annotation_position, options)
           return false unless File.exist?(file_name)
           old_content = File.read(file_name)
 
+          parser_klass = FileToParserMapper.map(file_name)
+
           begin
-            parsed_file = FileParser::ParsedFile.new(old_content, annotation, options).parse
+            parsed_file = FileParser::ParsedFile.new(old_content, annotation, parser_klass, options).parse
           rescue FileParser::AnnotationFinder::MalformedAnnotation => e
             warn "Unable to process #{file_name}: #{e.message}"
             warn "\t" + e.backtrace.join("\n\t") if @options[:trace]
@@ -39,11 +41,11 @@ def call(file_name, annotation, annotation_position, options)
           abort "AnnotateRb error. #{file_name} needs to be updated, but annotaterb was run with `--frozen`." if options[:frozen]
 
           updated_file_content = if !parsed_file.has_annotations?
-            AnnotatedFile::Generator.new(old_content, annotation, annotation_position, options).generate
+            AnnotatedFile::Generator.new(old_content, annotation, annotation_position, parser_klass, parsed_file, options).generate
           elsif options[:force]
-            AnnotatedFile::Generator.new(old_content, annotation, annotation_position, options).generate
+            AnnotatedFile::Generator.new(old_content, annotation, annotation_position, parser_klass, parsed_file, options).generate
           else
-            AnnotatedFile::Updater.new(old_content, annotation, annotation_position, options).update
+            AnnotatedFile::Updater.new(old_content, annotation, annotation_position, parsed_file, options).update
           end
 
           File.open(file_name, "wb") { |f| f.puts updated_file_content }

spec/lib/annotate_rb/model_annotator/annotated_file/generator_spec.rb

@@ -9,6 +9,8 @@
         file_content,
         new_annotations,
         annotation_position,
+        parser_klass,
+        parsed_file,
         options
       ]
     end
@@ -30,6 +32,10 @@ class User < ApplicationRecord
       ANNOTATIONS
     end
     let(:annotation_position) { :position_in_class }
+    let(:parser_klass) { AnnotateRb::ModelAnnotator::FileParser::CustomParser }
+    let(:parsed_file) do
+      AnnotateRb::ModelAnnotator::FileParser::ParsedFile.new(file_content, new_annotations, parser_klass, options).parse
+    end
 
     context 'when position is "before"' do
       let(:options) { AnnotateRb::Options.new({position_in_class: "before"}) }

spec/lib/annotate_rb/model_annotator/annotated_file/updater_spec.rb

@@ -9,11 +9,16 @@
         file_content,
         new_annotations,
         annotation_position,
+        parsed_file,
         options
       ]
     end
 
     let(:annotation_position) { :position_in_class }
+    let(:parsed_file) do
+      parser_klass = AnnotateRb::ModelAnnotator::FileParser::CustomParser
+      AnnotateRb::ModelAnnotator::FileParser::ParsedFile.new(file_content, new_annotations, parser_klass, options).parse
+    end
 
     context "with a foreign key constraint change" do
       let(:file_content) do