commands/wtf: add the --code flag

Fixes #1876 (Make `wtf` show code)

CHANGELOG.md

@@ -13,6 +13,8 @@
   ([#2024](https://github.com/pry/pry/pull/2024))
 * Added the `rc_file` config option that tells Pry the path to `pryrc`
   ([#2027](https://github.com/pry/pry/pull/2027))
+* Added the `--code` flag to the `wtf` command, which shows code for each
+  backtrace frame ([#2037](https://github.com/pry/pry/pull/2037))
 
 #### API changes
 

lib/pry/commands/wtf.rb

@@ -22,8 +22,11 @@ class Pry
         wtf -v
       BANNER
 
+      RUBY_FRAME_PATTERN = /\A(?<file>(.+)):(?<line>(\d+))/.freeze
+
       def options(opt)
         opt.on :v, :verbose, "Show the full backtrace"
+        opt.on :c, :code, "Show code corresponding to the backtrace frame"
       end
 
       def process
@@ -59,14 +62,31 @@ class Pry
       end
 
       def format_backtrace(backtrace)
-        return with_line_numbers(backtrace) if opts.verbose?
+        lines = trim_backtrace(backtrace).map do |frame|
+          next frame unless opts.code?
 
-        size_of_backtrace = [captures[0].size, 0.5].max * 10
-        with_line_numbers(backtrace.first(size_of_backtrace))
+          match = frame.match(RUBY_FRAME_PATTERN)
+          code = read_line(match[:file], match[:line].to_i)
+          [bold(frame), code].join("\n")
+        end
+
+        Pry::Code.new(lines.compact, 0, :text).with_line_numbers.to_s
       end
 
-      def with_line_numbers(backtrace)
-        Pry::Code.new(backtrace, 0, :text).with_line_numbers.to_s
+      def trim_backtrace(backtrace)
+        return backtrace if opts.verbose?
+
+        size_of_backtrace = [captures[0].size, 0.5].max * 10
+        backtrace.first(size_of_backtrace)
+      end
+
+      def read_line(file, line)
+        File.open(file, 'r') do |f|
+          (line - 1).times { f.gets }
+          f.gets
+        end
+      rescue Errno::ENOENT
+        nil
       end
     end
 

spec/commands/wtf_spec.rb

@@ -160,5 +160,56 @@ RSpec.describe Pry::Command::Wtf do
         end
       end
     end
+
+    context "when the code flag is present" do
+      let(:exception) do
+        error = RuntimeError.new('oops')
+        error.set_backtrace(
+          Array.new(6) { "#{__FILE__}:#{__LINE__}:in `<main>'" }
+        )
+        error
+      end
+
+      before do
+        expect(subject.opts).to receive(:code?).at_least(:once).and_return(true)
+      end
+
+      it "prints lines of code that exception frame references" do
+        subject.process
+        expect(subject.output.string).to eq(
+          "\e[1mException:\e[0m RuntimeError: oops\n" \
+          "--\n" \
+          "0: \e[1m#{__FILE__}:168:in `<main>'\e[0m\n" \
+          "             Array.new(6) { \"\#{__FILE__}:\#{__LINE__}:in `<main>'\" }\n" \
+          "1: \e[1m#{__FILE__}:168:in `<main>'\e[0m\n" \
+          "             Array.new(6) { \"\#{__FILE__}:\#{__LINE__}:in `<main>'\" }\n" \
+          "2: \e[1m#{__FILE__}:168:in `<main>'\e[0m\n" \
+          "             Array.new(6) { \"\#{__FILE__}:\#{__LINE__}:in `<main>'\" }\n" \
+          "3: \e[1m#{__FILE__}:168:in `<main>'\e[0m\n" \
+          "             Array.new(6) { \"\#{__FILE__}:\#{__LINE__}:in `<main>'\" }\n" \
+          "4: \e[1m#{__FILE__}:168:in `<main>'\e[0m\n" \
+          "             Array.new(6) { \"\#{__FILE__}:\#{__LINE__}:in `<main>'\" }\n"
+        )
+      end
+
+      context "and when referenced frame doesn't exist" do
+        before do
+          expect(File).to receive(:open).at_least(:once).and_raise(Errno::ENOENT)
+        end
+
+        it "skips code and prints only the backtrace frame" do
+          subject.process
+          expect(subject.output.string).to eq(
+            "\e[1mException:\e[0m RuntimeError: oops\n" \
+            "--\n" \
+            "0: \e[1m#{__FILE__}:168:in `<main>'\e[0m\n" \
+            "1: \e[1m#{__FILE__}:168:in `<main>'\e[0m\n" \
+            "2: \e[1m#{__FILE__}:168:in `<main>'\e[0m\n" \
+            "3: \e[1m#{__FILE__}:168:in `<main>'\e[0m\n" \
+            "4: \e[1m#{__FILE__}:168:in `<main>'\e[0m\n"
+          )
+        end
+      end
+    end
   end
 end