From e9c2f383eb6a9e3a4134f51f60ddf9c29d6ce165 Mon Sep 17 00:00:00 2001
From: Yorick Peterse <yorickpeterse@gmail.com>
Date: Wed, 5 Oct 2011 19:04:44 +0200
Subject: [PATCH] Pry now indents code similar to IRB.

Code is indented using the class Pry::Indent. This class uses an internal stack
that contains the indentation levels for each line of code. Certain keywords
such as "class" or "def" will add data to this stack so that the next line is
indented, other keywords such as "end" or "}" will remove data from the stack,
causing the next line to be un-indented.

Pry::Indent is hooked into Pry#retrieve_line as well as Pry#readline. This means
that both input strings as well as the ones displayed by "show-method" are
indented. Sadly due to the way Readline works input strings are indented similar
to IRB. This means that instead of the following:

    > class User
    >   def initialize
    >   end
    > end

You'll get the following:

    > class User
    >   def initialize
    >     end
    >   end

While annoying there doesn't seem to be a way to work around this issue. Luckily
the "show-method" command indents your code properly.

By default indentation is turned on. This can be turned off (or back on) using
the configuration item Pry.config.indent. However, if you turn this option off
after a method is defined "show-method" will still show it with indentation as
indentation happens on input rather than only when code is displayed.

For more information see Pry::Indent#indent in lib/pry/indent.rb.

Signed-off-by: Yorick Peterse <yorickpeterse@gmail.com>
---
 .gitignore              |   5 ++
 lib/pry.rb              |   1 +
 lib/pry/config.rb       |   4 ++
 lib/pry/indent.rb       | 136 ++++++++++++++++++++++++++++++++++++++++
 lib/pry/pry_class.rb    |   1 +
 lib/pry/pry_instance.rb |  27 +++++---
 test/helper.rb          |   5 +-
 test/test_indent.rb     | 127 +++++++++++++++++++++++++++++++++++++
 8 files changed, 297 insertions(+), 9 deletions(-)
 create mode 100644 lib/pry/indent.rb
 create mode 100644 test/test_indent.rb

diff --git a/.gitignore b/.gitignore
index d0b8649e..b4777dcf 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,3 +8,8 @@ pkg/
 coverage/
 .yardoc/
 /tags
+
+# Not sure if Pry devs use RVM or want the Gemfile.lock in the repo, thus I'm
+# ignoring them for now.
+.rvmrc
+Gemfile.lock
diff --git a/lib/pry.rb b/lib/pry.rb
index d364b110..fcd6ef7d 100644
--- a/lib/pry.rb
+++ b/lib/pry.rb
@@ -191,3 +191,4 @@ require "pry/plugins"
 require "pry/core_extensions"
 require "pry/pry_class"
 require "pry/pry_instance"
+require "pry/indent"
diff --git a/lib/pry/config.rb b/lib/pry/config.rb
index 37e236de..f5999214 100644
--- a/lib/pry/config.rb
+++ b/lib/pry/config.rb
@@ -126,6 +126,10 @@ class Pry
     #   The proc is passed the pry output object, the command string
     #   to eval, and a reference to the pry instance
     attr_accessor :system
+
+    # @return [TrueClass|FalseClass] Whether or not code should be indented
+    #  using Pry::Indent.
+    attr_accessor :indent
   end
 end
 
diff --git a/lib/pry/indent.rb b/lib/pry/indent.rb
new file mode 100644
index 00000000..2d33508a
--- /dev/null
+++ b/lib/pry/indent.rb
@@ -0,0 +1,136 @@
+require 'coderay'
+
+class Pry
+  ##
+  # Pry::Indent is a class that can be used to indent a number of lines
+  # containing Ruby code similar as to how IRB does it (but better). The class
+  # works by tokenizing a string using CodeRay and then looping over those
+  # tokens. Based on the tokens in a line of code that line (or the next one)
+  # will be indented or un-indented by 2 spaces.
+  #
+  # @author Yorick Peterse
+  # @since  04-10-2011
+  #
+  class Indent
+    # Array containing all the indentation levels.
+    attr_reader :stack
+
+    # The amount of spaces to insert for each indent level.
+    Spaces = '  '.freeze
+
+    # Array containing all the tokens that should increase the indentation
+    # level.
+    OpenTokens = [
+      'def',
+      'class',
+      'module',
+      '[',
+      '{',
+      'do',
+      'if',
+      'while',
+      'for'
+    ]
+
+    # Collection of tokens that decrease the indentation level.
+    ClosingTokens = ['end', ']', '}']
+
+    # Collection of token types that should be ignored. Without this list
+    # keywords such as "class" inside strings would cause the code to be
+    # indented incorrectly.
+    IgnoreTokens = [:space, :content, :string, :delimiter]
+
+    # Collection of tokens that should only increase the indentation level of
+    # the next line.
+    OpenTokensNext = ['else', 'elsif']
+
+    ##
+    # Creates a new instance of the class and starts with a fresh stack. The
+    # stack is used to keep track of the indentation level for each line of
+    # code.
+    #
+    # @author Yorick Peterse
+    # @since  05-10-2011
+    #
+    def initialize
+      @stack = []
+    end
+
+    ##
+    # Indents a string and returns it. This string can either be a single line
+    # or multiple ones.
+    #
+    # @example
+    #  str = <<TXT
+    #  class User
+    #  attr_accessor :name
+    #  end
+    #  TXT
+    #
+    #  # This would result in the following being displayed:
+    #  #
+    #  # class User
+    #  #   attr_accessor :name
+    #  # end
+    #  #
+    #  puts Pry::Indent.new.indent(str)
+    #
+    # @author Yorick Peterse
+    # @since  05-10-2011
+    # @param  [String] input The input string to indent.
+    # @return [String] The indented version of +input+.
+    #
+    def indent(input)
+      output = ''
+
+      input.lines.each do |line|
+        # Remove manually added indentation.
+        line   = line.strip + "\n"
+        tokens = CodeRay.scan(line, :ruby)
+
+        unless @stack.empty?
+          line = @stack[-1] + line
+        end
+
+        tokens.each_slice(2) do |token, kind|
+          next if IgnoreTokens.include?(kind)
+
+          if OpenTokensNext.include?(token) and @stack[-1]
+            line.sub!(@stack[-1], '')
+            break
+          # Start token found (such as "class"). Update the stack and indent the
+          # current line.
+          elsif OpenTokens.include?(token)
+            add  = ''
+            last = @stack[-1]
+
+            # Determine the amount of spaces to add to the stack and the line.
+            unless last.nil?
+              add = Spaces + last
+            end
+
+            # Don't forget to update the current line.
+            if @stack.empty?
+              line  = add + line
+              add  += Spaces
+            end
+
+            @stack.push(add)
+            break
+          # Stop token found. Remove the last number of spaces from the stack
+          # and un-indent the current line.
+          elsif ClosingTokens.include?(token)
+            @stack.pop
+
+            line = ( @stack[-1] || '' ) + line.strip + "\n"
+            break
+          end
+        end
+
+        output += line
+      end
+
+      return output.gsub!(/\s+$/, '')
+    end
+  end # Indent
+end # Pry
diff --git a/lib/pry/pry_class.rb b/lib/pry/pry_class.rb
index 3830f5fb..05bc8666 100644
--- a/lib/pry/pry_class.rb
+++ b/lib/pry/pry_class.rb
@@ -192,6 +192,7 @@ class Pry
     config.should_load_rc = true
     config.disable_auto_reload = false
     config.command_prefix = ""
+    config.indent = true
 
     config.plugins ||= OpenStruct.new
     config.plugins.enabled = true
diff --git a/lib/pry/pry_instance.rb b/lib/pry/pry_instance.rb
index 333fc5e1..b22b34e6 100644
--- a/lib/pry/pry_instance.rb
+++ b/lib/pry/pry_instance.rb
@@ -22,7 +22,6 @@ class Pry
   attr_reader :input_array
   attr_reader :output_array
 
-
   # Create a new `Pry` object.
   # @param [Hash] options The optional configuration parameters.
   # @option options [#readline] :input The object to use for input.
@@ -37,6 +36,7 @@ class Pry
 
     @command_processor = CommandProcessor.new(self)
     @binding_stack     = []
+    @indent            = Pry::Indent.new
   end
 
   # Refresh the Pry instance settings from the Pry class.
@@ -278,7 +278,9 @@ class Pry
   end
 
   # Read a line of input and check for ^d, also determine prompt to use.
-  # This method should not need to be invoked directly.
+  # This method should not need to be invoked directly. This method
+  # automatically indents the input value using Pry::Indent.
+  #
   # @param [String] eval_string The cumulative lines of input.
   # @param [Binding] target The target of the session.
   # @return [String] The line received.
@@ -286,20 +288,27 @@ class Pry
     current_prompt = select_prompt(eval_string.empty?, target.eval('self'))
     val = readline(current_prompt)
 
-    # exit session if we receive EOF character (^D)
+    # exit session if receive EOF character (^D)
     if !val
       output.puts ""
       Pry.config.control_d_handler.call(eval_string, self)
-      ""
+
+      return ""
     else
 
       # Change the eval_string into the input encoding (Issue 284)
-      # TODO: This wouldn't be necessary if the eval_string was constructed from input strings only.
-      if eval_string.empty? && val.respond_to?(:encoding) && val.encoding != eval_string.encoding
+      # TODO: This wouldn't be necessary if the eval_string was constructed from
+      # input strings only.
+      if eval_string.empty? && val.respond_to?(:encoding) \
+      && val.encoding != eval_string.encoding
         eval_string.force_encoding(val.encoding)
       end
 
-      val
+      if Pry.config.indent === true
+        val = @indent.indent(val)
+      end
+
+      return val
     end
   end
 
@@ -419,6 +428,10 @@ class Pry
   # @param [String] current_prompt The prompt to use for input.
   # @return [String] The next line of input.
   def readline(current_prompt="> ")
+    if Pry.config.indent === true
+      current_prompt += @indent.stack[-1] || ''
+    end
+
     handle_read_errors do
       if input == Readline
         line = input.readline(current_prompt, false)
diff --git a/test/helper.rb b/test/helper.rb
index 1bf5f89e..876aa12f 100644
--- a/test/helper.rb
+++ b/test/helper.rb
@@ -20,10 +20,11 @@ class << Pry
 
     Pry.color = false
     Pry.pager = false
-    Pry.config.should_load_rc = false
-    Pry.config.plugins.enabled = false
+    Pry.config.should_load_rc      = false
+    Pry.config.plugins.enabled     = false
     Pry.config.history.should_load = false
     Pry.config.history.should_save = false
+    Pry.config.indent              = false
     Pry.config.hooks = { }
   end
 end
diff --git a/test/test_indent.rb b/test/test_indent.rb
new file mode 100644
index 00000000..a1676e04
--- /dev/null
+++ b/test/test_indent.rb
@@ -0,0 +1,127 @@
+require File.expand_path('../helper', __FILE__)
+
+# Please keep in mind that any hash signs ("#") in the heredoc strings are
+# placed on purpose. Without these editors might remove the whitespace on empty
+# lines.
+describe Pry::Indent do
+  before do
+    @indent = Pry::Indent.new
+  end
+
+  it 'should indent an array' do
+    input  = "array = [\n10,\n15\n]"
+    output = "array = [\n  10,\n  15\n]"
+
+    @indent.indent(input).should === output
+  end
+
+  it 'should indent a hash' do
+    input  = "hash = {\n:name => 'Ruby'\n}"
+    output = "hash = {\n  :name => 'Ruby'\n}"
+
+    @indent.indent(input).should === output
+  end
+
+  it 'should indent a function' do
+    input  = "def\nreturn 10\nend"
+    output = "def\n  return 10\nend"
+
+    @indent.indent(input).should === output
+  end
+
+  it 'should indent a module and class' do
+    input        = "module Foo\n# Hello world\nend"
+    output       = "module Foo\n  # Hello world\nend"
+    input_class  = "class Foo\n# Hello world\nend"
+    output_class = "class Foo\n  # Hello world\nend"
+
+    @indent.indent(input).should       === output
+    @indent.indent(input_class).should === output_class
+  end
+
+  it 'should indent separate lines' do
+    @indent.indent('def foo').should   === 'def foo'
+    @indent.indent('return 10').should === '  return 10'
+    @indent.indent('end').should       === 'end'
+  end
+
+  it 'should properly indent nested code' do
+    input = <<TXT.strip
+module A
+module B
+class C
+attr_accessor :test
+# keep
+def number
+return 10
+end
+end
+end
+end
+TXT
+
+    output = <<TXT.strip
+module A
+  module B
+    class C
+      attr_accessor :test
+      # keep
+      def number
+        return 10
+      end
+    end
+  end
+end
+TXT
+
+    @indent.indent(input).should === output
+  end
+
+  it 'should indent statements such as if, else, etc' do
+    input = <<TXT.strip
+if a === 10
+#
+elsif a === 15
+#
+else
+#
+end
+#
+while true
+#
+end
+#
+for num in [10, 15, 20]
+#
+end
+#
+for num in [10, 15, 20] do
+#
+end
+TXT
+
+    output = <<TXT.strip
+if a === 10
+  #
+elsif a === 15
+  #
+else
+  #
+end
+#
+while true
+  #
+end
+#
+for num in [10, 15, 20]
+  #
+end
+#
+for num in [10, 15, 20] do
+  #
+end
+TXT
+
+    @indent.indent(input).should === output
+  end
+end