Add better support for activerecord relations and mongoid criteria

lib/draper.rb

@@ -23,6 +23,7 @@
 require 'draper/view_context'
 require 'draper/collection_decorator'
 require 'draper/undecorate'
+require 'draper/relation_decorator'
 require 'draper/decorates_assigned'
 require 'draper/railtie' if defined?(Rails)
 

lib/draper/decoratable.rb

@@ -49,15 +49,15 @@ def decorated?
 
     module ClassMethods
 
-      # Decorates a collection of objects. Used at the end of a scope chain.
+      # Decorates an ActiveRecord relation. Used at any point of the scope chain.
       #
       # @example
-      #   Product.popular.decorate
+      #   Product.popular.decorate.page(2)
       # @param [Hash] options
-      #   see {Decorator.decorate_collection}.
+      #   see {Decorator.decorate_relation}.
       def decorate(options = {})
-        collection = Rails::VERSION::MAJOR >= 4 ? all : scoped
-        decorator_class.decorate_collection(collection, options.reverse_merge(with: nil))
+        relation = Rails::VERSION::MAJOR >= 4 ? all : scoped
+        decorator_class.decorate_relation(relation, options.reverse_merge(with: nil))
       end
 
       def decorator_class?

lib/draper/decorator.rb

@@ -145,6 +145,23 @@ def self.decorate_collection(object, options = {})
       collection_decorator_class.new(object, options.reverse_merge(with: self))
     end
 
+    # Decorates an ActiveRecord relation. The class of the relation decorator
+    # is inferred from the decorator class if possible (e.g. `ProductDecorator`
+    # maps to `ProductsDecorator`), but otherwise defaults to
+    # {Draper::RelationDecorator}.
+    #
+    # @param [ActiveRecord::Relation] relation
+    #   relation to decorate.
+    # @option options [Class, nil] :with (self)
+    #   the decorator class used to decorate each item. When `nil`, it is
+    #   inferred from each item.
+    # @option options [Hash] :context
+    #   extra data to be stored in the collection decorator.
+    def self.decorate_relation(relation, options = {})
+      options.assert_valid_keys(:with, :context)
+      relation_decorator_class.new(relation, options.reverse_merge(with: self))
+    end
+
     # @return [Array<Class>] the list of decorators that have been applied to
     #   the object.
     def applied_decorators
@@ -247,6 +264,15 @@ def self.collection_decorator_class
       Draper::CollectionDecorator
     end
 
+    # @return [Class] the class created by {decorate_relation}.
+    def self.relation_decorator_class
+      name = collection_decorator_name
+      name.constantize
+    rescue NameError => error
+      raise if name && !error.missing_name?(name)
+      Draper::RelationDecorator
+    end
+
     private
 
     def self.inherited(subclass)

lib/draper/factory.rb

@@ -49,6 +49,7 @@ def call(options)
 
       def decorator
         return decorator_method(decorator_class) if decorator_class
+        return decorator_method(Draper::RelationDecorator) if relation?
         return object_decorator if decoratable?
         return decorator_method(Draper::CollectionDecorator) if collection?
         raise Draper::UninferrableDecoratorError.new(object.class)
@@ -59,21 +60,29 @@ def decorator
       attr_reader :decorator_class, :object
 
       def object_decorator
-        if collection?
+        if relation?
+          ->(object, options) { object.decorator_class.decorate_relation(object, options.reverse_merge(with: nil))}
+        elsif collection?
           ->(object, options) { object.decorator_class.decorate_collection(object, options.reverse_merge(with: nil))}
         else
           ->(object, options) { object.decorate(options) }
         end
       end
 
       def decorator_method(klass)
-        if collection? && klass.respond_to?(:decorate_collection)
+        if relation? && klass.respond_to?(:decorate_relation)
+          klass.method(:decorate_relation)
+        elsif collection? && klass.respond_to?(:decorate_collection)
           klass.method(:decorate_collection)
         else
           klass.method(:decorate)
         end
       end
 
+      def relation?
+        (defined?(ActiveRecord) && object.is_a?(ActiveRecord::Relation)) || (defined?(Mongoid) && object.is_a?(Mongoid::Criteria))
+      end
+
       def collection?
         object.respond_to?(:first) && !object.is_a?(Struct)
       end

lib/draper/relation_decorator.rb

@@ -0,0 +1,87 @@
+module Draper
+  class RelationDecorator
+    include Draper::ViewHelpers
+    extend Draper::Delegation
+
+    # @return [Class] the decorator class used to decorate this relation, as set by
+    #   {#initialize}.
+    attr_reader :decorator_class
+
+    # @return [Hash] extra data to be used in user-defined methods, and passed
+    #   to each item's decorator.
+    attr_accessor :context
+
+    # @param [ActiveRecord::Relation] relation
+    #   relation to decorate.
+    # @option options [Class, nil] :with (nil)
+    #   the decorator class used to decorate each item. When `nil`, each item's
+    #   {Decoratable#decorate decorate} method will be used.
+    # @option options [Hash] :context ({})
+    #   extra data to be stored in the relation decorator and used in
+    #   user-defined methods, and passed to each item's decorator.
+    def initialize(relation, options = {})
+      options.assert_valid_keys(:with, :context)
+      @relation = relation
+      @decorator_class = options[:with]
+      @decorator_class ||= klass.decorator_class if relation.respond_to?(:klass) && klass.respond_to?(:decorator_class)
+      @context = options.fetch(:context, {})
+    end
+
+    class << self
+      alias_method :decorate, :new
+    end
+
+    def to_ary
+      to_a
+    end
+
+    def to_s
+      "#<#{self.class.name} of #{decorator_class || "inferred decorators"} for #{relation.inspect}>"
+    end
+
+    def context=(value)
+      @context = value
+    end
+
+    # @return [true]
+    def decorated?
+      true
+    end
+
+    alias_method :decorated_with?, :instance_of?
+
+    def decorating_class
+      return decorator_class if decorator_class
+      self.class
+    end
+
+    def method_missing(method, *args, &block)
+      block ?
+        relation.send(method, *args, &proxy_block(&block)) :
+        handle_result(relation.send(method, *args))
+    end
+
+    def proxy_block(&original_block)
+      lambda { |data| original_block.call(handle_result(data)) }
+    end
+
+    def handle_result(result)
+      if (defined?(ActiveRecord) && result.is_a?(ActiveRecord::Relation)) ||
+        (defined?(Mongoid) && result.is_a?(Mongoid::Criteria))
+        return self.class.decorate(result, context: context)
+      elsif result.is_a?(Array)
+        return Decorator.collection_decorator_class.new(result, context: context)
+      elsif relation.respond_to?(:klass) && result.is_a?(relation.klass) && klass.respond_to?(:decorate)
+        return result.decorate(context: context)
+      else
+        return result
+      end
+    end
+
+    protected
+
+    # @return the relation being decorated.
+    attr_reader :relation
+
+  end
+end

spec/draper/decoratable_spec.rb

@@ -130,19 +130,19 @@ module Draper
     describe ".decorate" do
       let(:scoping_method) { Rails::VERSION::MAJOR >= 4 ? :all : :scoped }
 
-      it "calls #decorate_collection on .decorator_class" do
+      it "calls #decorate_relation on .decorator_class" do
         scoped = [Product.new]
         Product.stub scoping_method => scoped
 
-        Product.decorator_class.should_receive(:decorate_collection).with(scoped, with: nil).and_return(:decorated_collection)
-        expect(Product.decorate).to be :decorated_collection
+        Product.decorator_class.should_receive(:decorate_relation).with(scoped, with: nil).and_return(:decorated_relation)
+        expect(Product.decorate).to be :decorated_relation
       end
 
       it "accepts options" do
         options = {with: ProductDecorator, context: {some: "context"}}
         Product.stub scoping_method => []
 
-        Product.decorator_class.should_receive(:decorate_collection).with([], options)
+        Product.decorator_class.should_receive(:decorate_relation).with([], options)
         Product.decorate(options)
       end
     end

spec/draper/relation_decorator_spec.rb

@@ -0,0 +1,112 @@
+require 'spec_helper'
+require 'support/shared_examples/view_helpers'
+
+module Draper
+  describe RelationDecorator do
+    it_behaves_like "view helpers", RelationDecorator.new([])
+
+    describe "#initialize" do
+      describe "options validation" do
+
+        it "does not raise error on valid options" do
+          valid_options = {with: Decorator, context: {}}
+          expect{RelationDecorator.new(ActiveRecord::Relation.new, valid_options)}.not_to raise_error
+        end
+
+        it "raises error on invalid options" do
+          expect{RelationDecorator.new(ActiveRecord::Relation.new, foo: "bar")}.to raise_error ArgumentError, /Unknown key/
+        end
+      end
+    end
+
+    context "with context" do
+      it "stores the context itself" do
+        context = {some: "context"}
+        decorator = RelationDecorator.new(ActiveRecord::Relation.new, context: context)
+
+        expect(decorator.context).to be context
+      end
+    end
+
+    describe "#context=" do
+      it "updates the stored context" do
+        decorator = RelationDecorator.new(ActiveRecord::Relation.new, context: {some: "context"})
+        new_context = {other: "context"}
+
+        decorator.context = new_context
+        expect(decorator.context).to be new_context
+      end
+    end
+
+    it "returns a relation decorator when a scope is called on the decorated relation" do
+      module ActiveRecord
+        class Relation
+          include Draper::Decoratable
+          def some_scope; self ;end
+        end
+      end
+
+      klass = Product
+      klass.class_eval { def self.some_scope ; ActiveRecord::Relation.new ; end }
+      expect(Product).to respond_to(:some_scope)
+      proxy = RelationDecorator.new(klass)
+      expect(proxy.some_scope).to be_instance_of(proxy.class)
+    end
+
+    it 'supports chaining multiple scopes' do
+      module ActiveRecord
+        class Relation
+          include Draper::Decoratable
+          def some_scope; self ;end
+        end
+      end
+
+      klass = Product
+      klass.class_eval { def self.some_scope ; ActiveRecord::Relation.new ; end }
+      proxy = RelationDecorator.new(klass)
+      expect(proxy.some_scope.some_scope.some_scope).to be_instance_of(proxy.class)
+      expect(proxy.some_scope.some_scope.some_scope).to be_decorated
+    end
+
+    it 'supports converting the scope to an array of decorated objects' do
+      module ActiveRecord
+        class Relation
+          include Draper::Decoratable
+          def some_scope; self ;end
+
+          def to_a
+            [Product.new]
+          end
+        end
+      end
+
+      klass = Product
+      klass.class_eval { def self.some_scope ; ActiveRecord::Relation.new ; end }
+      proxy = RelationDecorator.new(klass)
+      expect(proxy.some_scope.to_ary).to be_a(Array)
+      expect(proxy.some_scope.to_ary.first).to be_instance_of(klass)
+      expect(proxy.some_scope.to_ary.first).to be_decorated
+    end
+
+    describe '#decorated?' do
+      it 'returns true' do
+        klass = Product
+        klass.class_eval { def self.some_scope ; ActiveRecord::Relation.new ; end }
+        decorator = ProductsRelationDecorator.new(Product.some_scope)
+
+        expect(decorator).to be_decorated
+      end
+    end
+
+    describe '#decorated_with?' do
+      it "checks if a decorator has been applied to a collection" do
+        klass = Product
+        klass.class_eval { def self.some_scope ; ActiveRecord::Relation.new ; end }
+        decorator = ProductsRelationDecorator.new(Product.some_scope)
+
+        expect(decorator).to be_decorated_with ProductsRelationDecorator
+        expect(decorator).not_to be_decorated_with OtherDecorator
+      end
+    end
+  end
+end

spec/dummy/app/models/post.rb

@@ -1,3 +1,4 @@
 class Post < ActiveRecord::Base
   # attr_accessible :title, :body
+  scope :active, ->{ where('1 = 1') }
 end

spec/dummy/spec/models/mongoid_post_spec.rb

@@ -4,5 +4,28 @@
 if defined?(Mongoid)
   describe MongoidPost do
     it_behaves_like "a decoratable model"
+
+    it 'correctly decorates on top of the criteria' do
+      MongoidPost.create
+      relation = MongoidPost.limit(1).decorate
+      expect(relation).to be_decorated_with Draper::RelationDecorator
+      expect(relation.first).to be_decorated_with MongoidPostDecorator
+      expect(relation.first).to be_a(MongoidPost)
+    end
+
+    it 'also supports interchanging scope order' do
+      MongoidPost.create
+      relation = MongoidPost.decorate.limit(1)
+      expect(relation).to be_decorated_with Draper::RelationDecorator
+      expect(relation.first).to be_decorated_with MongoidPostDecorator
+      expect(relation.first).to be_a(MongoidPost)
+    end
+
+    it 'works with in_groups_of' do
+      3.times { MongoidPost.create }
+      MongoidPost.decorate.in_groups_of(3, false) do |group|
+        expect(group.first).to be_decorated_with MongoidPostDecorator
+      end
+    end
   end
 end

spec/dummy/spec/models/post_spec.rb

@@ -3,4 +3,27 @@
 
 describe Post do
   it_behaves_like "a decoratable model"
+
+  it 'correctly decorates on top of the scopes' do
+    Post.create
+    relation = Post.limit(1).decorate
+    expect(relation).to be_decorated_with Draper::RelationDecorator
+    expect(relation.first).to be_decorated_with PostDecorator
+    expect(relation.first).to be_a(Post)
+  end
+
+  it 'also supports interchanging scope order' do
+    Post.create
+    relation = Post.decorate.limit(1)
+    expect(relation).to be_decorated_with Draper::RelationDecorator
+    expect(relation.first).to be_decorated_with PostDecorator
+    expect(relation.first).to be_a(Post)
+  end
+
+  it 'works with in_groups_of' do
+    3.times { Post.create }
+    Post.decorate.in_groups_of(3, false) do |group|
+      expect(group.first).to be_decorated_with PostDecorator
+    end
+  end
 end