Java Graphviz UML Generator

java_uml.rb

#!/usr/bin/env ruby

raise "Usage: ruby main.rb ../java_program/src" unless ARGV[0]
java_files = Dir.glob(ARGV[0] + "/**/*.java")

@klasses = []
java_files.each do |java_file|
  @class_name = nil
  @methods = []
  File.open(java_file).each_line do |line|
    
    # public class PictureHeuristic implements Heuristic<Picture> {
    class_name = line.match(/^public\ (class|abstract class|interface)\ (\w*)/)
    if class_name
      @class_name = class_name
    end
    
    # public double getHEstimate(Picture current) {
    method_name = line.match(/^\s*(public|public abstract|public static)[ ]?([|\w|\<|\>]*) (\w*)\(([\w|,| |\[|\]]*)\)/)
    if method_name
      @methods << method_name
    end
    
  end

  @klasses << [@class_name, @methods]
  
end

File.open("uml.gv", 'w') do |f|
  f.write "digraph uml_diagram {\n"
  f.write "graph[overlap=false, splines=true]\n"
  
  # Write classes
  @klasses.each do |klass|
    f.write "\"#{klass[0][2]}\" [shape=record, label=\"{#{klass[0][2]}|\\\n"

    # Methods
    klass[1].each do |method|
      f.write "#{method[3]}(#{method[4]})\\l\\\n"
    end
    f.write "}\"]\n"
  end
  
  # Write relations
  # TODO
  
  f.write "}"
  
end