lib/depcheck/output/graph_output.rb in depcheck-0.3.0 vs lib/depcheck/output/graph_output.rb in depcheck-0.4.0

@@ -1,20 +1,20 @@
 module Depcheck
   class GraphOutput
 
-    def post(objs)
+    def post(objs, include)
 
       # installs graphviz if needed
       system 'brew install graphviz' unless graphviz_installed?
 
       # Check that this is in the user's PATH after installing
       unless graphviz_installed?
         fail "graphviz is not in the user's PATH, or it failed to install"
       end
 
       # generate graph description
-      graph = generate_graph_description(objs)
+      graph = generate_graph_description(objs, include)
 
       # create temporary graph dot file
       file_name = 'graph'
       File.open("#{file_name}.dot", 'w') do |f|
         f.write(graph)
@@ -25,14 +25,16 @@
 
       # remove temporary file
       File.delete("#{file_name}.dot")
     end
 
-    def generate_graph_description(objs)
+    def generate_graph_description(objs, include)
       nodes = []
       objs.each do |obj|
         obj.dependencies.each do |dep|
-          nodes << { source: obj.name, dep: dep }
+          if obj.name.match(/#{include}/) || dep.match(/#{include}/)
+            nodes << { source: obj.name, dep: dep }
+          end
         end
       end
 
       desc = "digraph dependencies {\n node [fontname=monospace, fontsize=9, shape=box, style=rounded]\n"
       nodes.each do |node|