root/lang/ruby/ymldot/lib/ymldot.rb @ 19118

require 'yaml'
require 'pp'

class Entity
  def initialize(name, dependent, columns = [])
    @name = name
    @dependent = dependent
    @columns = columns
    @foreignkeys = []
  end

  def dependent?
    @dependent
  end

  def to_dot
    res = ""
    res << %Q!"#{@name}" [#{@dependent? 'shape=Mrecord, ' : 'shape=record, '}label=\"{#{@name}|!
    @foreignkeys.each{|f| res << "#{f}\\l"}
    @columns.each{|c| res << "#{c}\\l"}
    res << '}"]'
  end

  attr_accessor :name, :dependent, :columns, :foreignkeys
end


class Ymldot
  attr_accessor :file_name

  def initialize(filepath)
    open(filepath) do |f|
      str = f.read
      @node = YAML.load(str)
    end
    @file_name = $1 if filepath[/(\w+).yml\z/]
    @entities = {}
    @one_relations = []
    @many_relations = []
    eval
  end

  def dot_generate
    code = ""
    code += <<"EOS"
digraph #{@file_name} {
#{add_2_tab(config_to_dot)}
#{add_2_tab(entities_to_dot)}
#{add_2_tab(relations_to_dot)}
}
EOS
  end

  private
  def add_2_tab(str)
    res = ""
    str.each_line{|s| res << "  " << s}
    res
  end

  def config_to_dot
    res = ""
    res << "graph[overlap=false, splines=true]\n"
    res << %Q!node [fontname="#{@config["font"]}"]! if @config["font"]
    res
  end

  def entities_to_dot
    res = ""
    s = []
    @entities.each_value{|e| s << e }
    s.sort{|a, b| a.name <=> b.name}.each{|e| res << e.to_dot << "\n"}
    res
  end

  def relations_to_dot
    res = ""
    @one_relations.sort{|a, b| a[:self].name <=> b[:self].name}.each do |r|
      res << %Q!"#{r[:self].name}" -> "#{r[:have].name}" [arrowtail=none arrowhead=dot headlabel="1" taillabel="1"]! << "\n"
    end
    @many_relations.sort{|a, b| a[:self].name <=> b[:self].name}.each do |r|
      res << %Q!"#{r[:self].name}" -> "#{r[:have].name}" [arrowtail=none arrowhead=dot headlabel="n" taillabel="1"]! << "\n"
    end
    res
  end

  def eval
    @config = @node["config"]? @node["config"] : {}
    @node["tables"].each {|t| eval_entity(t) }
    @node["tables"].each {|t| eval_relation(t) }
  end

  def eval_entity(table)
    columns = table["columns"]; columns ||= []
    @entities[table["name"]] ||= Entity.new(table["name"], table["dependent"], columns)
  end

  def eval_relation(table)
    foreignkeys = table["foreignkeys"]
    tname = table["name"]
    eval_relation_belongs_to(foreignkeys["belongs_to"], tname) if foreignkeys["belongs_to"]
    eval_relation_has_many(foreignkeys["has_many"], tname) if foreignkeys["has_many"]
    eval_relation_has_one(foreignkeys["has_one"], tname) if foreignkeys["has_one"]
    eval_relation_hmabt(foreignkeys["has_many_and_belongs_to"], tname) if foreignkeys["has_many_and_belongs_to"]
  end

  def eval_relation_belongs_to(keys, tname)
    keys.each do |rel|
      @entities[tname].foreignkeys << "#{rel}ID(FK)"
    end
  end

  def eval_relation_has_many(keys, tname)
    keys.each do |rel|
      @many_relations << {:self => @entities[tname], :have => @entities[rel]}
    end
  end

  def eval_relation_has_one(keys, tname)
    keys.each do |rel|
      @one_relations << {:self => @entities[tname], :have => @entities[rel]}
    end
  end

  def eval_relation_hmabt(keys, tname)
    keys.each do |rel|
      join_tname = "#{rel}_#{tname}"
      return if @entities.has_key? join_tname
      join_tname = "#{tname}_#{rel}"
      @entities[join_tname] ||= Entity.new(join_tname, true)
      @entities[join_tname].foreignkeys << "#{tname}ID(FK)"
      @entities[join_tname].foreignkeys << "#{rel}ID(FK)"
      @many_relations << {:self => @entities[tname], :have => @entities[join_tname]}
      @many_relations << {:self => @entities[rel], :have => @entities[join_tname]}
    end
  end
end