From 16f69ee07ce66a4dda0ba735369aff1e47e6aeff Mon Sep 17 00:00:00 2001
From: Eduard Tsech <edtsech@gmail.com>
Date: Fri, 16 Sep 2011 20:42:33 +0200
Subject: [PATCH] Ability to pass block to decorator method

---
 lib/draper/model_support.rb       | 1 +
 spec/draper/model_support_spec.rb | 5 +++++
 spec/samples/active_record.rb     | 2 +-
 spec/samples/product_decorator.rb | 4 ++++
 4 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/lib/draper/model_support.rb b/lib/draper/model_support.rb
index 4fb671a..292ee13 100644
--- a/lib/draper/model_support.rb
+++ b/lib/draper/model_support.rb
@@ -1,5 +1,6 @@
 module Draper::ModelSupport
   def decorator
     @decorator ||= "#{self.class.name}Decorator".constantize.decorate(self)
+    block_given? ? yield(@decorator) : @decorator
   end
 end
diff --git a/spec/draper/model_support_spec.rb b/spec/draper/model_support_spec.rb
index fb97264..4a317d8 100644
--- a/spec/draper/model_support_spec.rb
+++ b/spec/draper/model_support_spec.rb
@@ -6,5 +6,10 @@ describe Draper::ModelSupport do
   describe '#decorator' do
     its(:decorator) { should be_kind_of(ProductDecorator) }
     its(:decorator) { should be(subject.decorator) }
+
+    it 'should have abillity to pass block' do
+      a = Product.new.decorator { |d| d.awesome_title }
+      a.should eql "Awesome Title"
+    end
   end
 end
diff --git a/spec/samples/active_record.rb b/spec/samples/active_record.rb
index 9b561b6..fc6fc48 100644
--- a/spec/samples/active_record.rb
+++ b/spec/samples/active_record.rb
@@ -4,4 +4,4 @@ module ActiveRecord
       name
     end
   end
-end
\ No newline at end of file
+end
diff --git a/spec/samples/product_decorator.rb b/spec/samples/product_decorator.rb
index f731419..eb81b74 100644
--- a/spec/samples/product_decorator.rb
+++ b/spec/samples/product_decorator.rb
@@ -1,3 +1,7 @@
 class ProductDecorator < Draper::Base
   decorates :product
+
+  def awesome_title
+    "Awesome Title"
+  end
 end