Bajena · Bajena · Jan 11, 2020 · Jan 9, 2020 · Jan 6, 2020 · Bajena
diff --git a/README.md b/README.md
@@ -163,6 +163,24 @@ class BlogPostSerializer < BaseSerializer
 end
 ```
 
+#### Example 6: Lazy dig through relationships
+In additional to previous example you may want to make use of nested lazy relationship without rendering of any nested record.
+There is an `lazy_dig` method to be used for that:
+
+```ruby
+class AuthorSerializer < BaseSerializer
+  lazy_relationship :address
+end
+
+class BlogPostSerializer < BaseSerializer
+  lazy_relationship :author
+
+  attribute :author_address do
+    lazy_dig(:author, :address)&.full_address
+  end
+end
+```
+
 ## Performance comparison with vanilla AMS
 
 In general the bigger and more complex your serialized records hierarchy is and the more latency you have in your DB the more you'll benefit from using this gem. 

diff --git a/lib/ams_lazy_relationships/core.rb b/lib/ams_lazy_relationships/core.rb
@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "ams_lazy_relationships/core/lazy_relationship_method"
+require "ams_lazy_relationships/core/lazy_dig_method"
 require "ams_lazy_relationships/core/relationship_wrapper_methods"
 require "ams_lazy_relationships/core/evaluation"
 
@@ -18,6 +19,7 @@ def self.ams_version
 
   def self.included(klass)
     klass.send :extend, ClassMethods
+    klass.send :include, LazyDigMethod
     klass.send :prepend, Initializer
 
     klass.send(:define_relationship_wrapper_methods)
@@ -48,7 +50,7 @@ module Initializer
     def initialize(*)
       super
 
-      self.class.send(:load_all_lazy_relationships, object)
+      self.class.send(:init_all_lazy_relationships, object)
     end
   end
 end
diff --git a/lib/ams_lazy_relationships/core/evaluation.rb b/lib/ams_lazy_relationships/core/evaluation.rb
@@ -8,56 +8,84 @@ module Evaluation
     LAZY_NESTING_LEVELS = 3
     NESTING_START_LEVEL = 1
 
+    # Loads the lazy relationship
+    #
+    # @param relation_name [Symbol] relation name to be loaded
+    # @param object [Object] Lazy relationships will be loaded for this record.
+    def load_lazy_relationship(relation_name, object)
+      lrm = lazy_relationships[relation_name]
+      unless lrm
+        raise ArgumentError, "Undefined lazy '#{relation_name}' relationship for '#{name}' serializer"
+      end
+
+      # We need to evaluate the promise right before serializer tries
+      # to touch it. Otherwise the various side effects can happen:
+      # 1. AMS will attempt to serialize nil values with a specific V1 serializer
+      # 2. `lazy_association ? 'exists' : 'missing'` expression will always
+      #     equal to 'exists'
+      # 3. `lazy_association&.id` expression can raise NullPointer exception
+      #
+      # Calling `__sync` will evaluate the promise.
+      init_lazy_relationship(lrm, object).__sync
+    end
+
     # Recursively loads the tree of lazy relationships
     # The nesting is limited to 3 levels.
     #
     # @param object [Object] Lazy relationships will be loaded for this record.
     # @param level [Integer] Current nesting level
-    def load_all_lazy_relationships(object, level = NESTING_START_LEVEL)
+    def init_all_lazy_relationships(object, level = NESTING_START_LEVEL)
       return if level >= LAZY_NESTING_LEVELS
       return unless object
 
       return unless lazy_relationships
 
       lazy_relationships.each_value do |lrm|
-        load_lazy_relationship(lrm, object, level)
+        init_lazy_relationship(lrm, object, level)
       end
     end
 
     # @param lrm [LazyRelationshipMeta] relationship data
     # @param object [Object] Object to load the relationship for
     # @param level [Integer] Current nesting level
-    def load_lazy_relationship(lrm, object, level = NESTING_START_LEVEL)
+    def init_lazy_relationship(lrm, object, level = NESTING_START_LEVEL)
       load_for_object = if lrm.load_for.present?
                           object.public_send(lrm.load_for)
                         else
                           object
                         end
 
       lrm.loader.load(load_for_object) do |batch_records|
-        deep_load_for_yielded_records(
+        deep_init_for_yielded_records(
           batch_records,
           lrm,
           level
         )
       end
     end
 
-    def deep_load_for_yielded_records(batch_records, lrm, level)
+    def deep_init_for_yielded_records(batch_records, lrm, level)
       # There'll be no more nesting if there's no
       # reflection for this relationship. We can skip deeper lazy loading.
       return unless lrm.reflection
 
       Array.wrap(batch_records).each do |r|
-        deep_load_for_yielded_record(r, lrm, level)
+        deep_init_for_yielded_record(r, lrm, level)
       end
     end
 
-    def deep_load_for_yielded_record(batch_record, lrm, level)
-      serializer = serializer_for(batch_record, lrm.reflection.options)
+    def deep_init_for_yielded_record(batch_record, lrm, level)
+      serializer = lazy_serializer_for(batch_record, lrm: lrm)
       return unless serializer
 
-      serializer.send(:load_all_lazy_relationships, batch_record, level + 1)
+      serializer.send(:init_all_lazy_relationships, batch_record, level + 1)
+    end
+
+    def lazy_serializer_for(object, lrm: nil, relation_name: nil)
+      lrm ||= lazy_relationships[relation_name]
+      return unless lrm&.reflection
+
+      serializer_for(object, lrm.reflection.options)
     end
   end
 end
diff --git a/lib/ams_lazy_relationships/core/lazy_dig_method.rb b/lib/ams_lazy_relationships/core/lazy_dig_method.rb
@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+module AmsLazyRelationships::Core
+  # Provides `lazy_dig` as an instance method for serializers, in order to make
+  # possible to dig relationships in depth just like `Hash#dig` do, keeping the
+  # laziness and N+1-free evaluation.
+  module LazyDigMethod
+    # @param relation_names [Array<Symbol>] the sequence of relation names
+    #   to dig through.
+    # @return [ActiveRecord::Base, Array<ActiveRecord::Base>, nil] ActiveRecord
+    #   objects found by digging through the sequence of nested relationships.
+    #   Singular or plural nature of returned value depends from the
+    #   singular/plural nature of the chain of relation_names.
+    #
+    # @example
+    #   class AuthorSerializer < BaseSerializer
+    #     lazy_belongs_to :address
+    #     lazy_has_many :rewards
+    #   end
+    #
+    #   class BlogPostSerializer < BaseSerializer
+    #     lazy_belongs_to :author
+    #
+    #     attribute :author_address do
+    #       # returns single AR object or nil
+    #       lazy_dig(:author, :address)&.full_address
+    #     end
+    #
+    #     attribute :author_rewards do
+    #       # returns an array of AR objects
+    #       lazy_dig(:author, :rewards).map(&:description)
+    #     end
+    #   end
+    def lazy_dig(*relation_names)
+      relationships = {
+        multiple: false,
+        data: [{
+          serializer: self.class,
+          object: object
+        }]
+      }
+
+      relation_names.each do |relation_name|
+        lazy_dig_relationship!(relation_name, relationships)
+      end
+
+      objects = relationships[:data].map { |r| r[:object] }
+
+      relationships[:multiple] ? objects : objects.first
+    end
+
+    private
+
+    def lazy_dig_relationship!(relation_name, relationships)
+      relationships[:data].map! do |serializer:, object:|
+        next_objects = lazy_dig_next_objects!(relation_name, serializer, object)
+        next unless next_objects
+
+        relationships[:multiple] ||= next_objects.respond_to?(:to_ary)
+
+        lazy_dig_next_relationships!(relation_name, serializer, next_objects)
+      end
+
+      relationships[:data].flatten!
+      relationships[:data].compact!
+    end
+
+    def lazy_dig_next_objects!(relation_name, serializer, object)
+      serializer&.send(
+        :load_lazy_relationship,
+        relation_name,
+        object
+      )
+    end
+
+    def lazy_dig_next_relationships!(relation_name, serializer, next_objects)
+      Array.wrap(next_objects).map do |next_object|
+        next_serializer = serializer.send(
+          :lazy_serializer_for,
+          next_object,
+          relation_name: relation_name
+        )
+
+        {
+          serializer: next_serializer,
+          object: next_object
+        }
+      end
+    end
+  end
+end
diff --git a/lib/ams_lazy_relationships/core/lazy_relationship_method.rb b/lib/ams_lazy_relationships/core/lazy_relationship_method.rb
@@ -39,15 +39,7 @@ def lazy_relationship(name, loader: nil, load_for: nil)
       @lazy_relationships[name] = lrm
 
       define_method :"lazy_#{name}" do
-        # We need to evaluate the promise right before serializer tries
-        # to touch it. Otherwise the various side effects can happen:
-        # 1. AMS will attempt to serialize nil values with a specific V1 serializer
-        # 2. `lazy_association ? 'exists' : 'missing'` expression will always
-        #     equal to 'exists'
-        # 3. `lazy_association&.id` expression can raise NullPointer exception
-        #
-        # Calling `__sync` will evaluate the promise.
-        self.class.send(:load_lazy_relationship, lrm, object).__sync
+        self.class.send(:load_lazy_relationship, name, object)
       end
     end