From e79538756254d48553a33522f77fd472cf4570f3 Mon Sep 17 00:00:00 2001
From: Iain Beeston <iainbeeston@users.noreply.github.com>
Date: Thu, 26 Oct 2023 13:26:58 +0100
Subject: [PATCH] Added support for annotate_rendered_view_with_filenames

ActionView's erb template handler can annotate html with comments to indicate which file any given snippet of html comes from. This is very useful in development, but it isn't supported by better-html.

This adds equivalent support to better-html, using the same approach that ActionView takes, so it should work automatically for anyone switching from rails' default erb handler.
---
 lib/better_html/better_erb.rb                  | 14 ++++++++++++--
 lib/better_html/config.rb                      |  1 +
 lib/better_html/railtie.rb                     |  6 ++++++
 .../better_erb/implementation_test.rb          | 18 ++++++++++++++++--
 4 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/lib/better_html/better_erb.rb b/lib/better_html/better_erb.rb
index 9ff935b..5fc8ecb 100644
--- a/lib/better_html/better_erb.rb
+++ b/lib/better_html/better_erb.rb
@@ -52,10 +52,20 @@ def generate(template, source)
         klass ||= self.class.erb_implementation
 
         escape = self.class.escape_ignore_list.include?(template.type)
+
+        options = {
+          escape: escape,
+          trim: (self.class.erb_trim_mode == "-"),
+        }
+        if BetterHtml.config.annotate_rendered_view_with_filenames && template.format == :html
+          options[:preamble] = "@output_buffer.safe_append='<!-- BEGIN #{template.short_identifier} -->';"
+          options[:postamble] = "@output_buffer.safe_append='<!-- END #{template.short_identifier} -->';" \
+            "@output_buffer.to_s"
+        end
+
         generator = klass.new(
           erb,
-          escape: escape,
-          trim: (self.class.erb_trim_mode == "-")
+          **options
         )
         generator.validate! if generator.respond_to?(:validate!)
         generator.src
diff --git a/lib/better_html/config.rb b/lib/better_html/config.rb
index 1582536..eb74401 100644
--- a/lib/better_html/config.rb
+++ b/lib/better_html/config.rb
@@ -15,6 +15,7 @@ class Config
     property :template_exclusion_filter
     property :lodash_safe_javascript_expression, default: -> { [/\AJSON\.stringify\(/] }
     property :disable_parser_validation, default: false
+    property :annotate_rendered_view_with_filenames, default: false
 
     def javascript_attribute_name?(name)
       javascript_attribute_names.any? { |other| other === name.to_s } # rubocop:disable Style/CaseEquality
diff --git a/lib/better_html/railtie.rb b/lib/better_html/railtie.rb
index ebd583a..328a07a 100644
--- a/lib/better_html/railtie.rb
+++ b/lib/better_html/railtie.rb
@@ -7,5 +7,11 @@ class Railtie < Rails::Railtie
     initializer "better_html.better_erb.initialization" do
       BetterHtml::BetterErb.prepend!
     end
+
+    config.after_initialize do
+      ActiveSupport.on_load(:action_view) do
+        BetterHtml.config.annotate_rendered_view_with_filenames = ActionView::Base.annotate_rendered_view_with_filenames
+      end
+    end
   end
 end
diff --git a/test/better_html/better_erb/implementation_test.rb b/test/better_html/better_erb/implementation_test.rb
index 8c356b9..2b75bd5 100644
--- a/test/better_html/better_erb/implementation_test.rb
+++ b/test/better_html/better_erb/implementation_test.rb
@@ -366,6 +366,20 @@ class ImplementationTest < ActiveSupport::TestCase
         end
       end
 
+      if ActionView.version >= Gem::Version.new("6.1")
+        test "with ActionView 6.1 comments are added to show the filename when annotate_rendered_view_with_filenames=true" do
+          config = build_config(annotate_rendered_view_with_filenames: true)
+          assert_equal "<!-- BEGIN _better_test.html.erb --><foo>bar<foo><!-- END _better_test.html.erb -->",
+            render("<foo>bar<foo>", config: config, filename: "_better_test.html.erb")
+        end
+      else
+        test "with ActionView 6.0 annotate_rendered_view_with_filenames=true does not change the output" do
+          config = build_config(annotate_rendered_view_with_filenames: true)
+          assert_equal "<foo>bar<foo>",
+            render("<foo>bar<foo>", config: config, filename: "_better_test.html.erb")
+        end
+      end
+
       test "capture works as intended" do
         output = render(<<-HTML)
       <%- foo = capture do -%>
@@ -420,13 +434,13 @@ def build_config(**options)
         BetterHtml::Config.new(**options)
       end
 
-      def render(source, config: build_config, locals: {})
+      def render(source, config: build_config, locals: {}, filename: "test.html.erb")
         old_config = BetterHtml.config
         BetterHtml.config = config
 
         ActionView::Template.new(
           source,
-          "test.html.erb",
+          filename,
           ActionView::Template::Handlers::ERB.new,
           virtual_path: "partial",
           format: :html,