Don't leave wrapper when removing annotations

lib/annotate/annotate_models.rb

@@ -403,10 +403,11 @@ module AnnotateModels
       end
     end
 
-    def remove_annotation_of_file(file_name)
+    def remove_annotation_of_file(file_name, options={})
       if File.exist?(file_name)
         content = File.read(file_name)
-        content.sub!(PATTERN, '')
+        wrapper_open = options[:wrapper_open] ? "# #{options[:wrapper_open]}\n" : ""
+        content.sub!(/(#{wrapper_open})?#{PATTERN}/, '')
 
         File.open(file_name, 'wb') { |f| f.puts content }
 
@@ -623,13 +624,13 @@ module AnnotateModels
             model_name = klass.name.underscore
             table_name = klass.table_name
             model_file_name = file
-            deannotated_klass = true if remove_annotation_of_file(model_file_name)
+            deannotated_klass = true if remove_annotation_of_file(model_file_name, options)
 
             get_patterns(matched_types(options)).
               map { |f| resolve_filename(f, model_name, table_name) }.
               each do |f|
                 if File.exist?(f)
-                  remove_annotation_of_file(f)
+                  remove_annotation_of_file(f, options)
                   deannotated_klass = true
                 end
               end

spec/annotate/annotate_models_spec.rb

@@ -514,6 +514,55 @@ class Foo < ActiveRecord::Base
 end
       EOS
     end
+
+    it "should remove opening wrapper" do
+      path = create "opening_wrapper.rb", <<-EOS
+# wrapper
+# == Schema Information
+#
+# Table name: foo
+#
+#  id                  :integer         not null, primary key
+#  created_at          :datetime
+#  updated_at          :datetime
+#
+
+class Foo < ActiveRecord::Base
+end
+      EOS
+
+      AnnotateModels.remove_annotation_of_file(path, wrapper_open: 'wrapper')
+
+      expect(content(path)).to eq <<-EOS
+class Foo < ActiveRecord::Base
+end
+      EOS
+    end
+
+    it "should remove closing wrapper" do
+      path = create "closing_wrapper.rb", <<-EOS
+class Foo < ActiveRecord::Base
+end
+
+# == Schema Information
+#
+# Table name: foo
+#
+#  id                  :integer         not null, primary key
+#  created_at          :datetime
+#  updated_at          :datetime
+#
+# wrapper
+
+      EOS
+
+      AnnotateModels.remove_annotation_of_file(path, wrapper_close: 'wrapper')
+
+      expect(content(path)).to eq <<-EOS
+class Foo < ActiveRecord::Base
+end
+      EOS
+    end
   end
 
   describe '#resolve_filename' do