Merge pull request #1999 from pry/color-printer-specs

Improve ColorPrinter and its unit tests

2 files changed, 92 insertions, 51 deletions

diff --git a/lib/pry/color_printer.rb b/lib/pry/color_printer.rb
index db6be35e..b339351a 100644
--- a/lib/pry/color_printer.rb
+++ b/lib/pry/color_printer.rb
@@ -1,71 +1,57 @@
 require 'pp'
+require 'English'
 
-# PP subclass for streaming inspect output in color.
 class Pry
+  # PP subclass for streaming inspect output in color.
   class ColorPrinter < ::PP
     Pry::SyntaxHighlighter.overwrite_coderay_comment_token!
 
-    OBJ_COLOR = begin
-      code = Pry::SyntaxHighlighter.keyword_token_color
-      if code.start_with? "\e"
-        code
-      else
-        "\e[0m\e[0;#{code}m"
-      end
+    def self.pp(obj, output = $DEFAULT_OUTPUT, max_width = 79)
+      queue = ColorPrinter.new(output, max_width, "\n")
+      queue.guard_inspect_key { queue.pp(obj) }
+      queue.flush
+      output << "\n"
     end
 
-    def self.pp(obj, out = $DEFAULT_OUTPUT, width = 79, newline = "\n")
-      q = ColorPrinter.new(out, width, newline)
-      q.guard_inspect_key { q.pp obj }
-      q.flush
-      out << "\n"
-    end
+    def pp(object)
+      return super unless object.is_a?(String)
 
-    def text(str, width = str.length)
-      # Don't recolorize output with color [Issue #751]
-      if str.include?("\e[")
-        super "#{str}\e[0m", width
-      elsif str.start_with?('#<') || str == '=' || str == '>'
-        super highlight_object_literal(str), width
-      else
-        super(SyntaxHighlighter.highlight(str), width)
-      end
+      # Avoid calling Ruby 2.4+ String#pretty_print that prints multiline
+      # Strings prettier
+      text(object.inspect)
+    rescue StandardError => exception
+      raise if exception.is_a?(Pry::Pager::StopPaging)
+
+      text(highlight_object_literal(inspect_object(object)))
     end
 
-    def pp(obj)
-      if obj.is_a?(String)
-        # Avoid calling Ruby 2.4+ String#pretty_print that prints multiline
-        # Strings prettier
-        text(obj.inspect)
+    def text(str, max_width = str.length)
+      if str.include?("\e[")
+        super("#{str}\e[0m", max_width)
+      elsif str.start_with?('#<') || %w[= >].include?(str)
+        super(highlight_object_literal(str), max_width)
       else
-        super
-      end
-    rescue StandardError => e
-      raise if e.is_a? Pry::Pager::StopPaging
-
-      begin
-        str = obj.inspect
-      rescue StandardError
-        # Read the class name off of the singleton class to provide a default
-        # inspect.
-        singleton = class << obj; self; end
-        ancestors = Pry::Method.safe_send(singleton, :ancestors)
-        klass  = ancestors.reject { |k| k == singleton }.first
-        obj_id = begin
-                   obj.__id__.to_s(16)
-                 rescue StandardError
-                   0
-                 end
-        str    = "#<#{klass}:0x#{obj_id}>"
+        super(SyntaxHighlighter.highlight(str), max_width)
       end
-
-      text highlight_object_literal(str)
     end
 
     private
 
     def highlight_object_literal(object_literal)
-      "#{OBJ_COLOR}#{object_literal}\e[0m"
+      code = Pry::SyntaxHighlighter.keyword_token_color
+      obj_color = code.start_with?("\e") ? code : "\e[0m\e[0;#{code}m"
+      "#{obj_color}#{object_literal}\e[0m"
+    end
+
+    def inspect_object(object)
+      object.inspect
+    rescue StandardError
+      # Read the class name off of the singleton class to provide a default
+      # inspect.
+      singleton = class << object; self; end
+      ancestors = Pry::Method.safe_send(singleton, :ancestors)
+      klass = ancestors.find { |k| k != singleton }
+      "#<#{klass}:0x#{object.__id__.to_s(16)}>"
     end
   end
 end
diff --git a/spec/color_printer_spec.rb b/spec/color_printer_spec.rb
index 6c2e72cb..13b2acf6 100644
--- a/spec/color_printer_spec.rb
+++ b/spec/color_printer_spec.rb
@@ -11,7 +11,7 @@ RSpec.describe Pry::ColorPrinter do
         end
       end
 
-      it "prints a string" do
+      it "prints a string with a newline" do
         described_class.pp(healthy_class.new, output)
         expect(output.string).to eq("string\n")
       end
@@ -40,5 +40,60 @@ RSpec.describe Pry::ColorPrinter do
           .to match(/\A\e\[32m#<BasicObject:0x.+>\e\[0m\e\[0m\n\z/)
       end
     end
+
+    context "when #inspect returns an object literal" do
+      let(:klass) do
+        Class.new do
+          def inspect
+            '#<Object:0x00007fe86bab53c8>'
+          end
+        end
+      end
+
+      it "prints the object inspect" do
+        described_class.pp(klass.new, output)
+        expect(output.string).to eq("\e[32m#<Object:0x00007fe86bab53c8>\e[0m\n")
+      end
+
+      context "and when SyntaxHighlighter returns a token starting with '\e'" do
+        before do
+          expect(Pry::SyntaxHighlighter).to receive(:keyword_token_color)
+            .and_return("\e[32m")
+        end
+
+        it "prints the object as is" do
+          described_class.pp(klass.new, output)
+          expect(output.string).to eq("\e[32m#<Object:0x00007fe86bab53c8>\e[0m\n")
+        end
+      end
+
+      context "and when SyntaxHighlighter returns a token that doesn't start with '\e'" do
+        before do
+          expect(Pry::SyntaxHighlighter).to receive(:keyword_token_color)
+            .and_return('token')
+        end
+
+        it "prints the object with escape characters" do
+          described_class.pp(klass.new, output)
+          expect(output.string)
+            .to eq("\e[0m\e[0;tokenm#<Object:0x00007fe86bab53c8>\e[0m\n")
+        end
+      end
+    end
+
+    context "when #inspect raises Pry::Pager::StopPaging" do
+      let(:klass) do
+        Class.new do
+          def inspect
+            raise Pry::Pager::StopPaging
+          end
+        end
+      end
+
+      it "propagates the error" do
+        expect { described_class.pp(klass.new, output) }
+          .to raise_error(Pry::Pager::StopPaging)
+      end
+    end
   end
 end