From 79af8ce6d7f46e7e79180dbd6bef6681c1329bd9 Mon Sep 17 00:00:00 2001
From: Nobuyoshi Nakada <nobu@ruby-lang.org>
Date: Thu, 1 Apr 2021 21:09:32 +0900
Subject: [PATCH] [ruby/irb] Colorize `__END__` as keyword

https://github.com/ruby/irb/commit/9b84018311
---
 lib/irb/color.rb       | 5 ++++-
 test/irb/test_color.rb | 1 +
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/lib/irb/color.rb b/lib/irb/color.rb
index a054bb20f8..cfbb3cc668 100644
--- a/lib/irb/color.rb
+++ b/lib/irb/color.rb
@@ -64,6 +64,7 @@ module IRB # :nodoc:
         on_alias_error:     [[RED, REVERSE],          ALL],
         on_class_name_error:[[RED, REVERSE],          ALL],
         on_param_error:     [[RED, REVERSE],          ALL],
+        on___end__:         [[GREEN],                 ALL],
       }
     rescue NameError
       # Give up highlighting Ripper-incompatible older Ruby
@@ -120,6 +121,7 @@ module IRB # :nodoc:
         symbol_state = SymbolState.new
         colored = +''
         length = 0
+        end_seen = false
 
         scan(code, allow_last_error: !complete) do |token, str, expr|
           # IRB::ColorPrinter skips colorizing fragments with any invalid token
@@ -138,10 +140,11 @@ module IRB # :nodoc:
             end
           end
           length += str.bytesize
+          end_seen = true if token == :on___end__
         end
 
         # give up colorizing incomplete Ripper tokens
-        if length != code.bytesize
+        unless end_seen or length == code.bytesize
           return Reline::Unicode.escape_for_print(code)
         end
 
diff --git a/test/irb/test_color.rb b/test/irb/test_color.rb
index 9976008124..a28ae06117 100644
--- a/test/irb/test_color.rb
+++ b/test/irb/test_color.rb
@@ -66,6 +66,7 @@ module TestIRB
         "\t" => "\t", # not ^I
         "foo(*%W(bar))" => "foo(*#{RED}#{BOLD}%W(#{CLEAR}#{RED}bar#{CLEAR}#{RED}#{BOLD})#{CLEAR})",
         "$stdout" => "#{GREEN}#{BOLD}$stdout#{CLEAR}",
+        "__END__" => "#{GREEN}__END__#{CLEAR}",
       }
 
       # specific to Ruby 2.7+