Colorize IRB's inspect result

Closes: https://github.com/ruby/ruby/pull/2150

3 files changed, 39 insertions, 2 deletions

diff --git a/lib/irb/color.rb b/lib/irb/color.rb
index 150e9e5666..ccb9002c23 100644
--- a/lib/irb/color.rb
+++ b/lib/irb/color.rb
@@ -41,6 +41,19 @@ module IRB # :nodoc:
         $stdout.tty? && ENV.key?('TERM')
       end
 
+      def inspect_colorable?(obj)
+        case obj
+        when Hash
+          obj.all? { |k, v| inspect_colorable?(k) && inspect_colorable?(v) }
+        when Array
+          obj.all? { |o| inspect_colorable?(o) }
+        when String, Symbol, Integer, Float, FalseClass, TrueClass, NilClass
+          true
+        else
+          false
+        end
+      end
+
       def clear
         return '' unless colorable?
         "\e[#{CLEAR}m"
diff --git a/lib/irb/inspector.rb b/lib/irb/inspector.rb
index f6f76712b8..7d2278a1f2 100644
--- a/lib/irb/inspector.rb
+++ b/lib/irb/inspector.rb
@@ -106,12 +106,22 @@ module IRB # :nodoc:
   Inspector.def_inspector([false, :to_s, :raw]){|v| v.to_s}
   Inspector.def_inspector([true, :p, :inspect]){|v|
     begin
-      v.inspect
+      result = v.inspect
+      if Color.inspect_colorable?(v)
+        result = Color.colorize_code(result)
+      end
+      result
     rescue NoMethodError
       puts "(Object doesn't support #inspect)"
     end
   }
-  Inspector.def_inspector([:pp, :pretty_inspect], proc{require "pp"}){|v| v.pretty_inspect.chomp}
+  Inspector.def_inspector([:pp, :pretty_inspect], proc{require "pp"}){|v|
+    result = v.pretty_inspect.chomp
+    if Color.inspect_colorable?(v)
+      result = Color.colorize_code(result)
+    end
+    result
+  }
   Inspector.def_inspector([:yaml, :YAML], proc{require "yaml"}){|v|
     begin
       YAML.dump(v)
diff --git a/test/irb/test_color.rb b/test/irb/test_color.rb
index 0c41613f51..088611c645 100644
--- a/test/irb/test_color.rb
+++ b/test/irb/test_color.rb
@@ -26,5 +26,19 @@ module TestIRB
         assert_equal(result, IRB::Color.colorize_code(code))
       end
     end
+
+    def test_inspect_colorable
+      {
+        1 => true,
+        2.3 => true,
+        ['foo', :bar] => true,
+        { a: 4 } => true,
+        Object.new => false,
+        Struct.new(:a) => false,
+        Struct.new(:a).new(1) => false,
+      }.each do |object, result|
+        assert_equal(result, IRB::Color.inspect_colorable?(object))
+      end
+    end
   end
 end