lib/rom/repository/relation_proxy/wrap.rb in rom-repository-1.0.2 vs lib/rom/repository/relation_proxy/wrap.rb in rom-repository-1.1.0

@@ -13,20 +13,19 @@
         # @param [Hash] options
         #
         # @return [RelationProxy]
         #
         # @api public
-        def wrap(options)
-          wraps = options.map { |(name, (relation, keys))|
-            relation.wrapped(name, keys)
-          }
+        def wrap(*names, **options)
+          new_wraps = wraps_from_names(names) + wraps_from_options(options)
 
-          relation = wraps.reduce(self) { |a, e|
-            a.relation.for_wrap(e.meta.fetch(:keys), e.base_name.relation)
+          relation = new_wraps.reduce(self) { |a, e|
+            name = e.meta[:wrap_from_assoc] ? e.meta[:combine_name] : e.base_name.relation
+            a.relation.for_wrap(e.meta.fetch(:keys), name)
           }
 
-          __new__(relation, meta: { wraps: wraps })
+          __new__(relation, meta: { wraps: wraps + new_wraps })
         end
 
         # Shortcut to wrap parents
         #
         # @example
@@ -50,11 +49,29 @@
         # so that correct mapper can be generated
         #
         # @return [RelationProxy]
         #
         # @api private
-        def wrapped(name, keys)
-          with(name: name, meta: { keys: keys, wrap: true, combine_name: name })
+        def wrapped(name, keys, wrap_from_assoc = false)
+          with(
+            name: name,
+            meta: {
+              keys: keys, wrap_from_assoc: wrap_from_assoc, wrap: true, combine_name: name
+            }
+          )
+        end
+
+        # @api private
+        def wraps_from_options(options)
+          options.map { |(name, (relation, keys))| relation.wrapped(name, keys) }
+        end
+
+        # @api private
+        def wraps_from_names(names)
+          names.map { |name|
+            assoc = associations[name]
+            registry[assoc.target].wrapped(name, assoc.combine_keys(__registry__), true)
+          }
         end
       end
     end
   end
 end