advanced module methods now in Pry::WrappedModule

* WrappedModule#source_location, WrappedModule#source, analogous to Pry::Method

lib/pry/code.rb

@@ -128,55 +128,15 @@ class Pry
         new(meth.source, start_line, meth.source_type)
       end
 
-      # FIXME: shares too much code with `from_module`
-      #
-      # Retrieve the source location of a module. Return value is in same
-      # format as Method#source_location. If the source location
-      # cannot be found this method returns `nil`.
-      #
-      # @param [Module] mod The module (or class).
-      # @return [Array<String, Fixnum>] The source location of the
-      #   module (or class).
-      def module_source_location(mod)
-        mod_type_string = mod.class.to_s.downcase
-        file, line = find_module_method_source_location(mod)
-
-        raise Pry::CommandError, "Can't find #{mod_type_string} code" if !file.is_a?(String)
-        class_regex = /#{mod_type_string}\s*(\w*)(::)?#{mod.name.split(/::/).last}/
-
-        if file == Pry.eval_path
-          all_lines = Pry.line_buffer.drop(1)
-        else
-          all_lines = File.readlines(file)
-        end
-
-        search_lines = all_lines[0..(line - 2)]
-        idx = search_lines.rindex { |v| class_regex =~ v }
-
-       [file,  idx + 1]
-      rescue Pry::RescuableException
-        nil
-      end
-
-      # FIXME: shares too much code with `module_source_location`
-      #
       # Attempt to extract the source code for module (or class) `mod`.
       #
       # @param [Module, Class] mod The module (or class) of interest.
       # @return [Code]
-      def from_module(mod)
-        file, line = module_source_location(mod)
+      def from_module(mod, start_line=nil)
+        mod = Pry::WrappedModule(mod)
 
-        raise CommandError, "Could not locate source for #{mod}!" if file.nil?
-
-        if file == Pry.eval_path
-          all_lines = Pry.line_buffer.drop(1)
-        else
-          all_lines = File.readlines(file)
-        end
-
-        mod_code = retrieve_complete_expression_from(all_lines[(line - 1)..-1])
-        new(mod_code, line, :ruby)
+        _, start_line = mod.source_location || 1
+        new(mod.source, start_line, :ruby)
       end
 
       protected
@@ -209,21 +169,6 @@ class Pry
 
           type
         end
-
-        def find_module_method_source_location(klass)
-          ims = Pry::Method.all_from_class(klass, false) + Pry::Method.all_from_obj(klass, false)
-
-          file, line = ims.each do |m|
-            break m.source_location if m.source_location && !m.alias?
-          end
-
-          if file && RbxPath.is_core_path?(file)
-            file = RbxPath.convert_path_to_full(file)
-          end
-
-          [file, line]
-        end
-
     end
 
     attr_accessor :code_type

lib/pry/default_commands/introspection.rb

@@ -220,10 +220,12 @@ class Pry
 
         def process_module(name)
           klass = target.eval(name)
-          file_name, line = Code.module_source_location(klass)
+
+          mod = Pry::WrappedModule(klass)
+          file_name, line = mod.source_location
           set_file_and_dir_locals(file_name)
           output.puts "\n#{Pry::Helpers::Text.bold('From:')} #{file_name} @ line #{line}:\n\n"
-          Code.from_module(klass).with_line_numbers(use_line_numbers?)
+          Code.from_module(mod).with_line_numbers(use_line_numbers?)
         end
 
         def use_line_numbers?

lib/pry/wrapped_module.rb

@@ -1,4 +1,16 @@
 class Pry
+  class << self
+    # If the given object is a `Pry::WrappedModule`, return it unaltered. If it's
+    # anything else, return it wrapped in a `Pry::WrappedModule` instance.
+    def WrappedModule(obj)
+      if obj.is_a? Pry::WrappedModule
+        obj
+      else
+        Pry::WrappedModule.new(obj)
+      end
+    end
+  end
+
   class WrappedModule
 
     attr_reader :wrapped
@@ -10,6 +22,8 @@ class Pry
     def initialize(mod)
       raise ArgumentError, "Tried to initialize a WrappedModule with a non-module #{mod.inspect}" unless ::Module === mod
       @wrapped = mod
+      @host_file_lines = nil
+      @source = nil
     end
 
     # The prefix that would appear before methods defined on this class.
@@ -69,5 +83,66 @@ class Pry
     def respond_to?(method_name)
       super || wrapped.send(method_name, *args, &block)
     end
+
+    # Retrieve the source for the module.
+    def source
+      if @source
+        @source
+      else
+        file, line = source_location
+        raise CommandError, "Could not locate source for #{wrapped}!" if file.nil?
+
+        @source = strip_leading_whitespace(Pry::Code.retrieve_complete_expression_from(@host_file_lines[(line - 1)..-1]))
+      end
+    end
+
+    # Retrieve the source location of a module. Return value is in same
+    # format as Method#source_location. If the source location
+    # cannot be found this method returns `nil`.
+    #
+    # @param [Module] mod The module (or class).
+    # @return [Array<String, Fixnum>] The source location of the
+    #   module (or class).
+    def source_location
+      mod_type_string = wrapped.class.to_s.downcase
+      file, line = find_module_method_source_location
+
+      return nil if !file.is_a?(String)
+
+      class_regex = /#{mod_type_string}\s*(\w*)(::)?#{wrapped.name.split(/::/).last}/
+
+      if file == Pry.eval_path
+        @host_file_lines ||= Pry.line_buffer.drop(1)
+      else
+        @host_file_lines ||= File.readlines(file)
+      end
+
+      search_lines = @host_file_lines[0..(line - 2)]
+      idx = search_lines.rindex { |v| class_regex =~ v }
+
+      [file,  idx + 1]
+    rescue Pry::RescuableException
+      nil
+    end
+
+    private
+
+    def find_module_method_source_location
+      ims = Pry::Method.all_from_class(wrapped, false) + Pry::Method.all_from_obj(wrapped, false)
+
+      file, line = ims.each do |m|
+        break m.source_location if m.source_location && !m.alias?
+      end
+
+      if file && RbxPath.is_core_path?(file)
+        file = RbxPath.convert_path_to_full(file)
+      end
+
+      [file, line]
+    end
+
+    def strip_leading_whitespace(text)
+      Pry::Helpers::CommandHelpers.unindent(text)
+    end
   end
 end