Rails class diagram (e.g models) with graphviz

class_diagram.rake

# (c) Loic Nageleisen
# License: MIT

Node = Struct.new(:type, :name, :attributes)
Edge = Struct.new(:type, :from, :to, :name)

class DiagramGraph
  def initialize
    @nodes = []
    @edges = []
  end

  def <<(item)
    case item
    when Node
      @nodes << item
    when Edge
      @edges << item
    end
  end

  def to_dot
    dot_header <<
      @nodes.map { |n| dot_node(n) }.join <<
      @edges.map { |e| dot_edge(e) }.join <<
    dot_footer
  end

  private

  def dot_header
    options = {}
    options.merge!(overlap: :false, splines: :true)
    #options.merge!(overlap: :scale)
    #options.merge!(rankdir: :LR)
    options.merge!(fontname: 'Helvetica')

    "digraph models_diagram {\n" \
    "\tgraph[#{dot_options(options)}]\n"
  end

  def dot_footer
    "}\n"
  end

  def dot_node(node)
    options = {}
    options.merge!(shape: :record)

    unless node.attributes.nil?
      attributes = node.attributes.sort
    end
    options.merge!(label: quote("{#{node.name}|#{dot_attributes(attributes)}}"))

    "\t#{quote(node.name)} [#{dot_options(options)}]\n"
  end

  def dot_edge(edge)
    options = {}

    options.merge!(label: quote(edge.name)) unless edge.name.nil?

    case edge.type
    when :belongs_to
      options.merge!(arrowtail: :dot, arrowhead: :vee)
    when :has_one
      options.merge!(arrowtail: :odot, arrowhead: :dot)
    when :has_many
      options.merge!(arrowtail: :odot, arrowhead: :normal)
    when :has_and_belongs_to_many
      options.merge!(arrowtail: :odot, arrowhead: :vee)
    end

    options.merge!(dir: :both, color: quote('#000000'))

    "\t#{quote(edge.from)} -> #{quote(edge.to)} [#{dot_options(options)}]\n"
  end

  def dot_options(options)
    options.map { |k, v| [k.to_s, v.to_s].join('=') }.join(', ')
  end

  def dot_attributes(attributes)
    attributes.join('\l') << '\l'
  end

  def quote(str)
    %w(" ").join(str.to_s)
  end
end

task diagram: [:environment] do
  blacklist = [
    #'Whatever',
  ]

  Rails.application.eager_load!

  graph = DiagramGraph.new
  ActiveRecord::Base.descendants.each do |model|
    next if blacklist.include?(model.name)

    graph << Node.new(:model, model.name, model.columns.map(&:name))
    model.reflect_on_all_associations.each do |association|
      next if association.options[:polymorphic]

      other_model = association.klass
      type = association.macro
      label = association.foreign_key
      graph << Edge.new(type, model.name, other_model.name, label)
    end
  end

  puts graph.to_dot
  #Tempfile.new('models') do |tmp|
  #  tmp.write(graph.to_dot)
  #  tmp.flush
  #  data = `cat #{tmp.path} | dot -Tsvg`
  #  puts data
  #end
end