From 67116b31d7db517472fb13938c94444edf96beea Mon Sep 17 00:00:00 2001
From: Andrew Emelianenko <emelianenko.web@gmail.com>
Date: Fri, 27 May 2016 15:28:03 +0300
Subject: [PATCH] Optimization and Fixes

---
 lib/pluck_to_hash.rb         | 44 ++++++++++++++++++++----------------
 spec/pluck_to_hash_spec.rb   | 17 +++++---------
 spec/pluck_to_struct_spec.rb | 10 ++++----
 spec/spec_helper.rb          |  2 +-
 4 files changed, 35 insertions(+), 38 deletions(-)

diff --git a/lib/pluck_to_hash.rb b/lib/pluck_to_hash.rb
index 760e2df..8c0c3cd 100644
--- a/lib/pluck_to_hash.rb
+++ b/lib/pluck_to_hash.rb
@@ -6,38 +6,42 @@ module PluckToHash
   module ClassMethods
     def pluck_to_hash(*keys)
       block_given = block_given?
-      keys = column_names if keys.blank?
-      formatted_keys = format_keys(keys)
+      keys, formatted_keys = format_keys(keys)
+      keys_one = keys.size == 1
+
       pluck(*keys).map do |row|
-        row = [row] if keys.size == 1
-        value = HashWithIndifferentAccess[formatted_keys.zip(row)]
-        yield(value) if block_given
-        value
+        value = HashWithIndifferentAccess[formatted_keys.zip(keys_one ? [row] : row)]
+        block_given ? yield(value) : value
       end
     end
 
     def pluck_to_struct(*keys)
       block_given = block_given?
-      keys = column_names if keys.blank?
-      formatted_keys = format_keys(keys)
+      keys, formatted_keys = format_keys(keys)
+      keys_one = keys.size == 1
 
-      struct = Struct.new(*formatted_keys)
+      struct = Struct.new(*formatted_keys.map(&:to_sym))
       pluck(*keys).map do |row|
-        row = [row] if keys.size == 1
-        value = struct.new(*row)
-        yield(value) if block_given
-        value
+        value = keys_one ? struct.new(*[row]) : struct.new(*row)
+        block_given ? yield(value) : value
       end
     end
 
     def format_keys(keys)
-      keys.map do |k|
-        case k
-        when String
-          k.split(/ as /i)[-1].to_sym
-        when Symbol
-          k
-        end
+      if keys.blank?
+        [column_names, column_names]
+      else
+        [
+          keys,
+          keys.map do |k|
+            case k
+            when String
+              k.split(/ as /i)[-1]
+            when Symbol
+              k.to_s
+            end
+          end
+        ]
       end
     end
 
diff --git a/spec/pluck_to_hash_spec.rb b/spec/pluck_to_hash_spec.rb
index 175e926..bff169c 100644
--- a/spec/pluck_to_hash_spec.rb
+++ b/spec/pluck_to_hash_spec.rb
@@ -35,21 +35,18 @@
       end
     end
 
-
     context 'the model does not have the attribute specified' do
       it 'raises an error' do
         expect do
           TestModel.all.pluck_to_hash(:foo)
-        end.to raise_error
+        end.to raise_error(ActiveRecord::StatementInvalid)
       end
     end
 
     context 'no models exist for the given criteria' do
       it 'returns an empty relation' do
-        expect do
-          result = TestModel.where(id: -1).pluck_to_hash(:id)
-          expect(result).to be_empty
-        end.to_not raise_error
+        result = TestModel.where(id: -1).pluck_to_hash(:id)
+        expect(result).to be_empty
       end
     end
 
@@ -111,16 +108,14 @@
         it 'raises an error' do
           expect do
             TestModel.all.pluck_h(:foo)
-          end.to raise_error
+          end.to raise_error(ActiveRecord::StatementInvalid)
         end
       end
 
       context 'no models exist for the given criteria' do
         it 'returns an empty relation' do
-          expect do
-            result = TestModel.where(id: -1).pluck_h(:id)
-            expect(result).to be_empty
-          end.to_not raise_error
+          result = TestModel.where(id: -1).pluck_h(:id)
+          expect(result).to be_empty
         end
       end
 
diff --git a/spec/pluck_to_struct_spec.rb b/spec/pluck_to_struct_spec.rb
index 5075a4a..2557695 100644
--- a/spec/pluck_to_struct_spec.rb
+++ b/spec/pluck_to_struct_spec.rb
@@ -12,7 +12,7 @@
 
     it 'plucks the ids, test_attributes of the objects to a struct correctly' do
       TestModel.all.pluck_to_struct(:test_attribute, :id).each_with_index do |model, ix|
-        id = ix+1
+        id = ix + 1
         expect(model).to be_a(Struct)
         expect(model.test_attribute).to eq("test#{id}")
         expect(model.id).to eq(id)
@@ -30,16 +30,14 @@
       it 'raises an error' do
         expect do
           TestModel.all.pluck_h(:foo)
-        end.to raise_error
+        end.to raise_error(ActiveRecord::StatementInvalid)
       end
     end
 
     context 'no models exist for the given criteria' do
       it 'returns an empty relation' do
-        expect do
-          result = TestModel.where(id: -1).pluck_h(:id)
-          expect(result).to be_empty
-        end.to_not raise_error
+        result = TestModel.where(id: -1).pluck_h(:id)
+        expect(result).to be_empty
       end
     end
 
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index fc71db8..7e94062 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -17,4 +17,4 @@
 
 class TestModel < ActiveRecord::Base
   serialize :serialized_attribute, Array
-end
\ No newline at end of file
+end