type_balancer_rails 0.1.2 → 0.1.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b8cb44e1021bb6681468e80eb02a256ced78b07e7dfbe825075743f84041925a
-  data.tar.gz: a4b24000bd076ba214c300c3fb6df31205aa2e5c902b18b283c927d8b3b3a2a2
+  metadata.gz: 5ebd00e75d4c63af53f5be038f01f920eb6e65dcf3c91ae079a10b20435d62ce
+  data.tar.gz: 67d328f6ae917d27f88455dce176073dd8da0cf8dd7bb2be460c886c7a56086f
 SHA512:
-  metadata.gz: 66f4ae2af899d916990f755ecd458a3a65e5ac4b3f2f0f0018162ab4a8e26a6bfa68ef2af3b035ec82fcbab4b4118b4b46fccaebb31e719e7842ccbd4ff694ab
-  data.tar.gz: 8bac8d9abd629d29ca25a49aae129c7ec546842c18903b674b441311cf01e9627350e3cc48ae6ee0f7c31f1d6ed59d072437b1db749e4b37fe9ec0a604d6d245
+  metadata.gz: 4a82e84484c322cd4ee02c92950cf3235daa989b21097c845a5e41a454084afe883ee15444729992e7a63b59318508dbd7f89a15d1abe3d371f8de2cc1e1c551
+  data.tar.gz: e427b4224fbed0bd83dcda7dd2ce97683dc1c7df54a2fb673ce15b9d545ee029f87637fb5a8cff685cb7570f7d0e7aeaf3c82e0a5466906773326a2d9d794486

data/lib/type_balancer/rails/active_record_extension.rb

@@ -19,9 +19,15 @@ module TypeBalancer
       end
 
       class_methods do
-        def balance_by_type(options = {})
+        # Accepts either a symbol (type field) or options hash
+        def balance_by_type(type_field = nil, **options)
+          if type_field.is_a?(Hash)
+            options = type_field
+            type_field = options[:type_field]
+          end
+
           self.type_balancer_options = {
-            type_field: options[:type_field] || :type
+            type_field: type_field || options[:type_field] || :type
           }
 
           return [] unless respond_to?(:all)

data/lib/type_balancer/rails/collection_methods.rb

@@ -23,31 +23,43 @@ module TypeBalancer
         records = to_a
         return empty_relation if records.empty?
 
-        # Get type field from options or model configuration
         type_field = fetch_type_field(options)
-
-        # Balance records using TypeBalancer
-        balanced = TypeBalancer.balance(records, type_field: type_field)
+        balanced   = TypeBalancer.balance(records, type_field: type_field)
         return empty_relation if balanced.nil?
 
-        # Handle pagination if requested
-        if options[:page] || options[:per_page]
-          page = (options[:page] || 1).to_i
-          per_page = (options[:per_page] || 20).to_i
-          offset = (page - 1) * per_page
-          balanced = balanced[offset, per_page] || []
-        end
-
-        # Return as relation
-        TestRelation.new(balanced)
+        paged = apply_pagination(balanced, options)
+        build_result(paged)
       end
 
-      def all_records = to_a
-
       private
 
+      def apply_pagination(records, options)
+        return records unless options[:page] || options[:per_page]
+
+        page      = (options[:page] || 1).to_i
+        per_page  = (options[:per_page] || 20).to_i
+        offset    = (page - 1) * per_page
+        records[offset, per_page] || []
+      end
+
+      def build_result(balanced)
+        if klass.respond_to?(:where) && klass != TestModel
+          ids     = balanced.map(&:id)
+          results = klass.where(id: ids).to_a
+          results_by_id = results.index_by(&:id)
+          ordered = ids.map { |id| results_by_id[id] }
+          self.class.new(ordered)
+        else
+          self.class.new(balanced)
+        end
+      end
+
       def empty_relation
-        TestRelation.new([])
+        if klass.respond_to?(:none)
+          klass.none
+        else
+          []
+        end
       end
 
       def fetch_type_field(options)

data/lib/type_balancer/rails/version.rb

@@ -2,6 +2,6 @@
 
 module TypeBalancer
   module Rails
-    VERSION = '0.1.2'
+    VERSION = '0.1.3'
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: type_balancer_rails
 version: !ruby/object:Gem::Version
-  version: 0.1.2
+  version: 0.1.3
 platform: ruby
 authors:
 - Carl Smith
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-04-26 00:00:00.000000000 Z
+date: 2025-04-27 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord