Merge pull request #219 from voormedia/pr-189

support for 'only_models_include_depth' option

README.md

@@ -60,6 +60,7 @@ warn: true
 title: sample title
 exclude: null
 only: null
+only_recursion_depth: null
 prepend_primary: false
 cluster: false
 ```

examples/erdconfig.example

@@ -17,5 +17,6 @@ warn: true
 title: sample title
 exclude: null
 only: null
+only_recursion_depth: null
 prepend_primary: false
 cluster: false

lib/rails_erd.rb

@@ -51,8 +51,9 @@ def default_options
         :title, true,
         :exclude, nil,
         :only, nil,
+        :only_recursion_depth, nil,
         :prepend_primary, false,
-        :cluster, false
+        :cluster, false,
       ]
     end
   end

lib/rails_erd/cli.rb

@@ -49,6 +49,11 @@
     desc "Filter to only include listed models in diagram."
   end
 
+  option :only_recursion_depth do
+    long "--only_recursion_depth=INTEGER"
+    desc "Recurses into relations specified by --only upto a depth N."
+  end
+
   option :exclude do
     long "--exclude"
     desc "Filter to exclude listed models in diagram."

lib/rails_erd/diagram.rb

@@ -125,6 +125,14 @@ def create
     def generate
       instance_eval(&callbacks[:setup])
 
+      if options.only_recursion_depth.present?
+        depth = options.only_recursion_depth.to_i
+        options.only.dup.each do |class_name|
+          options.only += recurse_into_relationships(@domain.entity_by_name(class_name), depth)
+        end
+        options.only.uniq!
+      end
+
       filtered_entities.each do |entity|
         instance_exec entity, filtered_attributes(entity), &callbacks[:each_entity]
       end
@@ -138,6 +146,26 @@ def generate
       end
     end
 
+    def recurse_into_relationships(entity, max_level, current_level = 0)
+      return [] unless entity
+      return [] if max_level == current_level
+
+      relationships = entity.relationships.reject{|r| r.indirect? || r.recursive?}
+
+      relationships.map do |relationship|
+        other_entitiy = if relationship.source == entity
+                          relationship.destination
+                        else
+                          relationship.source
+                        end
+        if other_entitiy and !other_entitiy.generalized?
+          [other_entitiy.name] + recurse_into_relationships(other_entitiy, max_level, current_level + 1)
+        else
+          []
+        end
+      end.flatten.uniq
+    end
+
     def save
       instance_eval(&callbacks[:save])
     end