lib/serega/plugins/batch/lib/modules/attribute_normalizer.rb in serega-0.16.0 vs lib/serega/plugins/batch/lib/modules/attribute_normalizer.rb in serega-0.17.0

@@ -40,25 +40,40 @@
 
         def prepare_batch
           batch = init_opts[:batch]
           return unless batch
 
-          # take loader
-          loader = batch[:loader]
+          loader = prepare_batch_loader(batch[:loader])
 
-          # take key
           key = batch[:key] || self.class.serializer_class.config.batch.default_key
-          proc_key =
-            if key.is_a?(Symbol)
-              proc { |object| object.public_send(key) }
-            else
-              key
-            end
+          key = prepare_batch_key(key)
 
-          # take default value
           default = batch.fetch(:default) { many ? FROZEN_EMPTY_ARRAY : nil }
 
-          {loader: loader, key: proc_key, default: default}
+          {loader: loader, key: key, default: default}
+        end
+
+        def prepare_batch_key(key)
+          return proc { |object| object.public_send(key) } if key.is_a?(Symbol)
+
+          params_count = SeregaUtils::ParamsCount.call(key, max_count: 2)
+          case params_count
+          when 0 then proc { key.call }
+          when 1 then proc { |object| key.call(object) }
+          else key
+          end
+        end
+
+        def prepare_batch_loader(loader)
+          return loader if loader.is_a?(Symbol)
+
+          params_count = SeregaUtils::ParamsCount.call(loader, max_count: 3)
+          case params_count
+          when 0 then proc { loader.call }
+          when 1 then proc { |object| loader.call(object) }
+          when 2 then proc { |object, context| loader.call(object, context) }
+          else loader
+          end
         end
       end
     end
   end
 end