From 9208a850a10d9685aa31fea1e0556710805ee828 Mon Sep 17 00:00:00 2001
From: Ryan Fitzgerald <rwfitzge@gmail.com>
Date: Sat, 22 Sep 2012 17:34:09 -0700
Subject: [PATCH] Refactor AR show-model and test --grep

---
 lib/pry-rails/commands/show_models.rb | 63 +++++++++++++++++----------
 spec/show_models_spec.rb              | 11 ++++-
 2 files changed, 50 insertions(+), 24 deletions(-)

diff --git a/lib/pry-rails/commands/show_models.rb b/lib/pry-rails/commands/show_models.rb
index 243dfde..ca1fde2 100644
--- a/lib/pry-rails/commands/show_models.rb
+++ b/lib/pry-rails/commands/show_models.rb
@@ -17,21 +17,42 @@ PryRails::Commands.create_command "show-models", "Show all defined models." do
     Rails.application.eager_load!
 
     if defined?(ActiveRecord::Base)
-      models = ActiveRecord::Base.descendants.map do |mod|
-        model_string = mod.to_s + "\n"
-        if mod.table_exists?
-          model_string << mod.columns.map { |col| "  #{col.name}: #{col.type.to_s}" }.join("\n")
+      models = ActiveRecord::Base.descendants
+
+      display(models.map do |model|
+        model_string = model.to_s + "\n"
+
+        if model.table_exists?
+          model.columns.each do |column|
+            model_string << "  #{column.name}: #{column.type.to_s}\n"
+          end
         else
-          model_string << "  Table doesn't exist"
+          model_string << "  Table doesn't exist\n"
         end
-        mod.reflections.each do |model,ref|
-          model_string << "\n  #{ref.macro.to_s} #{model}"
-          model_string << " through #{ref.options[:through]}" unless ref.options[:through].nil?
+
+        model.reflections.each do |model, reflection|
+          model_string << "  #{reflection.macro.to_s} #{model}"
+
+          if reflection.options[:through].present?
+            model_string << " through #{reflection.options[:through]}\n"
+          else
+            model_string << "\n"
+          end
         end
+
         model_string
-      end.join("\n")
-    elsif defined?(Mongoid::Document)
-      models = get_files.map do |path|
+      end.join)
+    end
+
+    if defined?(Mongoid::Document)
+      models = []
+      ObjectSpace.each_object do |o|
+        if o.is_a?(Class) && o.ancestors.include?(Mongoid::Document)
+          models << o
+        end
+      end
+
+      display(models.map do |model|
         mod = extract_class_name(path)
         model_string = "\033[1;34m#{mod.to_s}\033[0m\n"
         begin
@@ -54,19 +75,17 @@ PryRails::Commands.create_command "show-models", "Show all defined models." do
         rescue Exception
           STDERR.puts "Warning: exception #{$!} raised while trying to load model class #{path}"
         end
-      end.join("\n")
+      end.join("\n"))
+    end
+  end
+
+  def display(string)
+    if opts.present?(:G)
+      regexp = Regexp.new(opts[:G], Regexp::IGNORECASE)
+      string = string.gsub(regexp) { |s| text.red(s) }
     end
 
-    models.gsub!(Regexp.new(opts[:G] || ".", Regexp::IGNORECASE)) { |s| text.red(s) } unless opts[:G].nil?
-    output.puts models
-  end
-
-  def get_files(prefix ='')
-    Dir.glob(prefix << "app/models/**/*.rb")
-  end
-
-  def extract_class_name(filename)
-    filename.split('/')[2..-1].collect { |i| i.camelize }.join('::').chomp(".rb")
+    output.puts string
   end
 
   def kind_of_relation(string)
diff --git a/spec/show_models_spec.rb b/spec/show_models_spec.rb
index 74eeba4..ae84e6e 100644
--- a/spec/show_models_spec.rb
+++ b/spec/show_models_spec.rb
@@ -2,8 +2,8 @@
 
 require 'spec_helper'
 
-describe 'show-models' do
-  it 'should print a list of ActiveRecord models' do
+describe "show-models" do
+  it "should print a list of ActiveRecord models" do
     mock_pry('show-models', 'exit-all').must_equal <<OUTPUT
 Beer
   id: integer
@@ -28,4 +28,11 @@ Pokemon
   has_many beers through hacker
 OUTPUT
   end
+
+  it "should highlight the given phrase with --grep" do
+    output = mock_pry('show-models --grep beer', 'exit-all')
+
+    output.must_include "\e[0;31mBeer\e[0m"
+    output.must_include "has_many \e[0;31mbeer\e[0ms through hacker"
+  end
 end