graph.rb

# Generates infrastructure graph in DOT format
# Usage:
#   $ find . -name '*.pp' | while read file; do cat $file; echo; done | \
#     ruby graph.rb > puppet.dot

NODE      = /^\s*(class|define|node)(\s+)([\w:]+)/
RELATION  = /^\s*(inherits|include)(\s+)([\w:]+)/
RELATION2 = /^\s*class(\s+)\{\s*(['"]?)([\w:]+)/
USAGE     = /^\s*([\w:]*[\w])(\s*)\{(\s*)/
IGNORE    = %w{file package service cron exec Exec line plugin user group class
               concat concat::fragment line stage}

$nodes = {}

while line = gets
  line.gsub!(/#.*/, '')

  if line =~ NODE
    current_node = "\"#{$3}\""
    $nodes[current_node] = {
      "type"     => $1,
      "uses"     => [],
      "include"  => [],
      "require"  => [],
      "inherits" => []}
  end

  if current_node
    if line =~ RELATION
      $nodes[current_node][$1] << "\"#{$3}\""
    end

    if line =~ RELATION2
      $nodes[current_node]["include"] << "\"#{$3}\""
    end

    if line =~ USAGE
      $nodes[current_node]["uses"] << "\"#{$1}\"" unless IGNORE.include?($1)
    end
  end
end

def aggregate(array)
  array.inject({}) do |agg,obj|
    agg.has_key?(obj) ? agg[obj] += 1 : agg[obj] = 1
    agg
  end
end

puts "digraph G {"
$nodes.keys.sort.each do |n|
  a = $nodes[n]

  puts "#{n}#{a['type'] == 'node' ? ' [shape=box]' : ''};"

  ag = aggregate(a['uses'])
  ag.keys.sort.each do |k|
    v = ag[k]
    puts "#{n} -> {#{k}} [style=dotted#{v>1 ? ",label=#{v}" : ""}];"
  end

  puts "#{n} -> {#{a['include'].sort.join('; ')}};" if a['include'].any?
  puts "#{n} -> {#{a['require'].sort.join('; ')}};" if a['require'].any?
  puts "#{n} -> {#{a['inherits'].sort.join('; ')}};" if a['inherits'].any?
end
puts "}"