railroad in railroad-0.3.0

- old
+ new
@@ -1,241 +1,433 @@
 #!/usr/bin/env ruby
 
 # RailRoad - RoR diagrams generator
-# http://railroad.rubyforge.org/
+# http://railroad.rubyforge.org
 #
+#   RailRoad generates models and controllers diagrams in DOT language
+#   for a Rails application.
+#
 # Copyright 2007 - Javier Smaldone (http://www.smaldone.com.ar)
 #
 # This program is free software; you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
 # the Free Software Foundation; either version 2 of the License, or
 # (at your option) any later version.
 #
-# Purpose:
-#   RailRoad generates models and controllers diagrams in DOT language
-#   for a Rails application.
 
-@app_name  = "railroad"
-@app_human_name = "RailRoad"
-@version   = "0.1.2"
-@copyright = "Copyright (C) 2007 Javier Smaldone"
+APP_NAME  = "railroad"
+APP_HUMAN_NAME = "RailRoad"
+APP_VERSION   = [0,3,0]
+COPYRIGHT = "Copyright (C) 2007 Javier Smaldone"
 
-def print_help
-  print "Usage: #{@app_name} [options] [command]\n\n"
-  print "Options:\n"
-  print "  -a, --all\tInclude all models (not ActiveRecord::Base derived)\n"
-  print "  -b, --brief\tGenerate compact diagram (no attributes)\n"
-  print "  -c, --compact\tConcentrate edges (only for modules diagram)\n"
-  print "  -l, --label\tAdd a label with diagram info (type, date, version)\n"
-  print "  -i, --inherit\tInclude inheritance relations\n"
-  print "  -h, --help\tDisplay this help message and exit\n"
-  print "  -v, --version\tDisplay version information and exit\n\n"
-  print "Commands:\n"
-  print "  -M\tGenerate models diagram\n"
-  print "  -C\tGenerate controllers diagram\n"
-  print "\n"
-end
 
-def node(name)
+class OptionsParser
+
+  require 'optparse'
+  require 'ostruct'
+
+  def self.parse(args)
+
+    options = OpenStruct.new
+
+    options.all = false
+    options.brief = false
+    options.inheritance = false
+    options.join = false
+    options.label = false
+    options.modules = false
+    options.hide_types = false
+    options.hide_public = false
+    options.hide_protected = false
+    options.hide_private = false
+    options.command = ''
+
+    opts = OptionParser.new do |opts|
+      opts.banner = "Usage: #{APP_NAME} [options] command"
+      opts.separator ""
+      opts.separator "Common options:"
+      opts.on("-b", "--brief", "Generate compact diagram", 
+              "  (no attributes nor methods)") do |b|
+        options.brief = b
+      end
+      opts.on("-i", "--inheritance", "Include inheritance relations") do |i|
+        options.inheritance = i
+      end
+      opts.on("-l", "--label", "Add a label with diagram information", 
+              "  (type, date, migration, version)") do |l|
+        options.label = l
+      end
+      opts.on("-o", "--output FILE", "Write diagram to file FILE") do |f|
+        options.output = f
+      end
+      opts.separator ""
+      opts.separator "Models diagram options:"
+      opts.on("-a", "--all", "Include all models", 
+              "  (not only ActiveRecord::Base derived)") do |a|
+        options.all = a
+      end
+      opts.on("--hide-types", "Hide attributes type") do |h|
+        options.hide_types = h
+      end
+      opts.on("-j", "--join", "Concentrate edges") do |j|
+        options.join = j
+      end
+      opts.on("-m", "--modules", "Include modules") do |m|
+        options.modules = m
+      end
+      opts.separator ""
+      opts.separator "Controllers diagram options:"
+      opts.on("--hide-public", "Hide public methods") do |h|
+        options.hide_public = h
+      end
+      opts.on("--hide-protected", "Hide protected methods") do |h|
+        options.hide_protected = h
+      end
+      opts.on("--hide-private", "Hide private methods") do |h|
+        options.hide_private = h
+      end
+      opts.separator ""
+      opts.separator "Other options:"
+      opts.on("-h", "--help", "Show this message") do
+        print "#{opts}\n"
+        exit
+      end
+      opts.on("--version", "Show version and copyright") do
+        print "#{APP_HUMAN_NAME} version #{APP_VERSION.join('.')}\n\n" +
+              "#{COPYRIGHT}\n" +
+              "This is free software; see the source for copying conditions.\n\n"
+        exit
+      end
+      opts.separator ""
+      opts.separator "Commands (you must supply one of these):"
+      opts.on("-M", "--models", "Generate models diagram") do |c|
+         if options.command == 'controllers'
+           $stderr.print "Error: Can't generate models AND controllers diagram\n\n"
+           exit 1
+         else 
+           options.command = 'models'        
+         end
+      end
+      opts.on("-C", "--controllers", "Generate controllers diagram") do |c|
+         if options.command == 'models'
+           $stderr.print "Error: Can't generate models AND controllers diagram\n\n"
+           exit 1
+         else 
+           options.command = 'controllers'        
+         end
+      end
+      opts.separator ""
+      opts.separator "For bug reporting instructions and additional information, please see:"
+      opts.separator "http://railroad.rubyforge.org"
+    end
+
+    begin
+      opts.parse!(args)
+    rescue OptionParser::AmbiguousOption
+      option_error "Ambiguous option", opts 
+    rescue OptionParser::InvalidOption
+      option_error "Invalid option", opts
+    rescue OptionParser::InvalidArgument
+      option_error "Invalid argument", opts
+    rescue OptionParser::MissingArgument
+      option_error "Missing argument", opts
+    rescue
+      option_error "Unknown error", opts
+    end
+    options
+  end  # self.parse
+ 
+  private 
+
+  def self.option_error(msg, opts)
+    $stderr.print "Error: #{msg}\n\n #{opts}\n"
+    exit 1
+  end
+
+end  # class OptionsParser
+
+
+class AppDiagram
+
+  def initialize(options)
+    @options = options
+    load_environment
+  end
+
+  private 
+
+  # Quotes a class name
+  def node(name)
     '"' + name + '"'
-end
+  end
 
-def print_info(type)
-  print "\t_diagram_info [shape=\"plaintext\", label=\"Diagram: #{type}\\l"
-  print "Date: #{Time.now.strftime "%b %d %Y - %H:%M"}\\l"
-  print "Migration version: #{ActiveRecord::Migrator.current_version}\\l"
-  print "Generated by #{@app_human_name} #{@version}\\l\""
-  print ", fontsize=14]\n"
-end
+  # Print diagram label
+  def print_info(type)
+    print "\t_diagram_info [shape=\"plaintext\", label=\"Diagram: #{type}\\l" +
+          "Date: #{Time.now.strftime "%b %d %Y - %H:%M"}\\l" + 
+          "Migration version: #{ActiveRecord::Migrator.current_version}\\l" +
+          "Generated by #{APP_HUMAN_NAME} #{APP_VERSION.join('.')}\\l\""+
+          ", fontsize=14]\n"
+  end
 
-def print_error(type)
-    STDERR.print "Error loading #{type}. (Are you running "
-    STDERR.print "#{@app_name} on the aplication's root directory?)\n\n"
-end
+  def print_error(type)
+    STDERR.print "Error loading #{type}.\n  (Are you running " +
+                 "#{APP_NAME} on the aplication's root directory?)\n\n"
+  end
 
-def generate_models
-  # Load model classes
-  begin
-    Dir.glob("app/models/**/*.rb") {|m| require m }
-  rescue LoadError
-    print_error "model classes"
-    raise
+  # Load Rails application's environment
+  def load_environment
+    begin
+      require "config/environment"
+    rescue LoadError
+      print_error "app environment"
+      raise
+    end
   end
 
-  # DOT diagram header
-  print "digraph models_diagram {\n"
-  print "\tgraph[concentrate=true]\n" if @compact
-  print_info "Models" if @label
+end # class AppDiagram
 
-  models_files = Dir.glob("app/models/**/*.rb")
-  models_files.each do |m|
-    # Extract the class name from the file name
-    current_class = m.split('/')[2..-1].join('/').split('.').first.camelize.constantize
 
+class ModelsDiagram < AppDiagram
+
+  # Generate models diagram
+  def generate
+
+    load_classes
+
+    # DOT diagram header
+    print "digraph models_diagram {\n" 
+    print "\tgraph[concentrate=true]\n" if @options.join
+    print_info "Models" if @options.label
+
+    # Processed habtm associations
+    @habtm = []
+
+    models_files = Dir.glob("app/models/**/*.rb")
+    models_files.each do |m|
+      # Extract the class name from the file name
+      current_class = m.split('/')[2..-1].join('/').split('.').first.camelize.constantize
+      process_class current_class  
+    end 
+    print "}\n"
+
+
+  end # generate
+
+  private
+
+  # Load model classes
+  def load_classes
+    begin
+      Dir.glob("app/models/**/*.rb") {|m| require m }
+    rescue LoadError
+      print_error "model classes"
+      raise
+    end
+  end  # load_classes
+
+  # Process model class
+  def process_class(current_class)
+
+    printed = false
+    
     # Is current_clas derived from ActiveRecord::Base?
     if current_class.respond_to?'reflect_on_all_associations'
 
       # Print the node
-      if @brief
+      if @options.brief
         print "\t#{node(current_class.name)}\n"
       else 
-        print "\t#{node(current_class.name)} "
-        print "[shape=Mrecord, label=\"{#{current_class.name}|"
+        print "\t#{node(current_class.name)} " +
+              "[shape=Mrecord, label=\"{#{current_class.name}|"
         current_class.content_columns.each do |a|
-          print "#{a.name} :#{a.type.to_s}\\l" 
+          print "#{a.name}"
+          print " :#{a.type.to_s}" unless @options.hide_types
+          print "\\l" 
         end
         print "}\"]\n"
       end
-
+      printed = true
       # Iterate over the class associations
-      current_class.reflect_on_all_associations.each do |association|
+      current_class.reflect_on_all_associations.each { |a|
+        process_association(current_class, a)
+      }
+    elsif @options.all && (current_class.is_a? Class)
+      # Not ActiveRecord::Base model
+      if @options.brief
+        print "\t#{node(current_class.name)} [shape=box]\n"
+      else 
+        print "\t#{node(current_class.name)}" +
+              "[shape=record, label=\"{#{current_class.name}|}\"]\n"
+      end
+      printed = true
+    elsif @options.modules && (current_class.is_a? Module)
+        print "\t#{node(current_class.name)}" +
+              "[shape=box, style=dotted, label=\"#{current_class.name}\"]\n"      
+    end
 
-        # Only non standard association names needs a label
-        if association.class_name != association.name.to_s.singularize.camelize
-          params = "label=\"#{association.name.to_s}\", "
-        else
-          params = ""
-        end
+    # Only consider meaningful inheritance relations for printed classes
+    if @options.inheritance && printed && 
+       (current_class.superclass != ActiveRecord::Base) &&
+       (current_class.superclass != Object)
+      print "\t\t#{node(current_class.name)} -> " +
+            "#{node(current_class.superclass.name)}" +
+            " [arrowhead=\"onormal\"]\n"
+    end      
 
-        # Skip "belongs_to" associations
-        if association.macro.to_s != 'belongs_to'
-          # Put a tail label with the association arity
-          if association.macro.to_s == 'has_one' 
-            params += 'taillabel="1"'
-          else
-            params += 'taillabel="n"'
-          end
+  end # process_class
 
-          # Print the edge
-          print "\t\t#{node(current_class.name)} -> "
-          print "#{node(association.class_name)} "
-          print "[ #{params} ]\n"
-        end
-      end # Associations
+  # Process model association
+  def process_association(current_class, association)
 
-      if @inherit && (current_class.superclass != ActiveRecord::Base)
-        # Print the inheritance edge
-        print "\t\t#{node(current_class.name)} -> "
-	print "#{node(current_class.superclass.name)}"
-        print " [arrowhead=\"onormal\"]\n"
-      end      
+    # Only non standard association names needs a label
+    if association.class_name != association.name.to_s.singularize.camelize
+      params = "label=\"#{association.name.to_s}\", "
+    else
+      params = ""
+    end
 
-    elsif @all_classes
-      # Not ActiveRecord::Base model
-      if @brief
-        print "\t#{node(current_class.name)} [shape=box]\n"
-      else 
-        print "\t#{node(current_class.name)}"
-        print "[shape=record, label=\"{#{current_class.name}|}\"]\n"
+    # Skip "belongs_to" associations
+    if association.macro.to_s != 'belongs_to'
+      # Put a tail label with the association arity
+      if association.macro.to_s == 'has_one' 
+        params += 'taillabel="1"'
+      else
+        params += 'taillabel="n"'
       end
-    end
-  end # Dir
-  print "}\n"
-end # generate_models
 
-def generate_controllers
-  # Load controller classes
-  begin
-    # ApplicationController must be loaded first
-    require "app/controllers/application.rb" 
-    Dir.glob("app/controllers/**/*_controller.rb") do |c| 
-      require c
+      # Use double-headed arrows for habtm and has_many, :through
+      if association.macro.to_s == 'has_and_belongs_to_many' ||
+         (association.macro.to_s == 'has_many' && association.options[:through])
+        if @habtm.include? "#{association.class_name} -> #{current_class.name}"
+	  # habtm association already considered
+          return
+        else
+          params += ', headlabel="n", arrowtail="normal"'
+          @habtm << "#{current_class.name} -> #{association.class_name}"
+        end
+      end  
+      # Print the edge
+      print "\t\t#{node(current_class.name)} -> " +
+            "#{node(association.class_name)} " +
+            "[ #{params} ]\n"
     end
-  rescue LoadError
-    print_error "controller classes"
-    raise
-  end
+  end # process_association
 
-  # DOT diagram header
-  print "digraph controllers_diagram {\n"
-  print "\tgraph[overlap=false, splines=true]\n"
-  print_info "Controllers" if @label
+end # class ModelsDiagram
 
-  app_controller = ApplicationController
 
-  controllers_files = Dir.glob("app/controllers/**/*_controller.rb")
-  controllers_files << 'app/controllers/application.rb'
-  controllers_files.each do |c|
-    # Extract the class name from the file name
-    class_name = c.split('/')[2..-1].join('/').split('.').first.camelize
+class ControllersDiagram < AppDiagram
 
-    # ApplicationController's file is 'application.rb'
-    if class_name == 'Application' 
-      class_name = 'ApplicationController' 
+  # Generate controllers diagram
+  def generate
+
+    load_classes
+
+    # DOT diagram header
+    print "digraph controllers_diagram {\n" +
+          "\tgraph[overlap=false, splines=true]\n"
+    print_info "Controllers" if @options.label
+
+    @app_controller = ApplicationController
+
+    controllers_files = Dir.glob("app/controllers/**/*_controller.rb")
+    controllers_files << 'app/controllers/application.rb'
+    controllers_files.each do |c|
+
+      # Extract the class name from the file name
+      class_name = c.split('/')[2..-1].join('/').split('.').first.camelize
+
+      # ApplicationController's file is 'application.rb'
+      if class_name == 'Application' 
+        class_name = 'ApplicationController' 
+      end 
+
+      process_class class_name.constantize
+
     end 
-    current_class = class_name.constantize
+    print "}\n"
+  end # generate
 
+  private
+
+  # Load controller classes
+  def load_classes
+    begin
+      # ApplicationController must be loaded first
+      require "app/controllers/application.rb" 
+      Dir.glob("app/controllers/**/*_controller.rb") do |c| 
+        require c
+      end
+    rescue LoadError
+      print_error "controller classes"
+      raise
+    end
+  end # load_classes
+
+  # Proccess controller class
+  def process_class(current_class)
     # Print the node
-    if @brief
+    if @options.brief
       print "\t#{node(current_class.name)}\n"
-    else 
-      print "\t#{node(current_class.name)} "
-      print "[shape=Mrecord, label=\"{#{current_class.name}|"
-      current_class.public_instance_methods(false).sort.each do |m|
+    elsif current_class.is_a? Class 
+      print "\t#{node(current_class.name)} " +
+            "[shape=Mrecord, label=\"{#{current_class.name}|"
+      current_class.public_instance_methods(false).sort.each { |m|
         print "#{m}\\l" 
-      end
-      print "|"
-      current_class.protected_instance_methods(false).sort.each do |m|
+      } unless @options.hide_public
+      print "|" 
+      current_class.protected_instance_methods(false).sort.each { |m|
         print "#{m}\\l" 
-      end
-      print "|"
-      current_class.private_instance_methods(false).sort.each do |m|
+      } unless @options.hide_protected
+      print "|" 
+      current_class.private_instance_methods(false).sort.each { |m|
         print "#{m}\\l" 
-      end
-      print "}\"]\n"
+      } unless @options.hide_private
+      print "}\"]\n" 
+    elsif @options.modules && (current_class.is_a? Module)
+        print "\t#{node(current_class.name)}" +
+              "[shape=box, style=dotted, label=\"#{current_class.name}\"]\n"      
     end
-    if @inherit && (app_controller.subclasses.include? current_class.name)
-      # Print the inheritance edge
-      print "\t#{node(current_class.name)} -> "
-      print "#{node(current_class.superclass.name)} "
-      print "[arrowhead=\"onormal\"]\n"
+
+    # Print the inheritance edge
+    if @options.inheritance && 
+       (@app_controller.subclasses.include? current_class.name)
+      print "\t#{node(current_class.name)} -> " +
+            "#{node(current_class.superclass.name)} " +
+            "[arrowhead=\"onormal\"]\n"
     end
-  end # Dir
-  print "}\n"
-end # generate_controllers
 
-# Main program 
-if ARGV.include?('--help') || ARGV.include?('-h')
-  print_help
-  exit
-elsif ARGV.include?('--version') || ARGV.include?('-v')
-  print "#{@app_name} #{@version}\n#{@copyright}\n\n"
-  exit
-end
+  end # process_class
 
-@all_classes = ARGV.include?('--all') || ARGV.include?('-a')
-@brief = ARGV.include?('--brief') || ARGV.include?('-b')
-@compact = ARGV.include?('--compact') || ARGV.include?('-c')
-@label = ARGV.include?('--label') || ARGV.include?('-l')
-@inherit = ARGV.include?('--inherit') || ARGV.include?('-i')
+end # class ControllersDiagram
 
-params = %w(-a --all -b --brief -c --compact -l --label -i --inherit -M -C)
-error_params = ARGV - params
-if ! error_params.empty? 
-  # Unknown parameter
-  STDERR.print "Error: wrong parameter(s) #{error_params}\n"
-  STDERR.print "(try '#{@app_name} --help')\n\n"
+
+# Main program
+
+options = OptionsParser.parse ARGV
+
+if options.command == 'models'
+  diagram = ModelsDiagram.new options
+elsif options.command == 'controllers'
+  diagram = ControllersDiagram.new options
+else
+  $stderr.print "Error: You must supply a command\n" +
+                "  (try #{APP_NAME} -h)\n\n"
   exit 1
 end
-if ARGV.include?('-M') && ARGV.include?('-C')
-  STDERR.print "Error: you can only specify one command\n"
-  STDERR.print "(try '#{@app_name} --help')\n\n"
-  exit 1
-elsif ! (ARGV.include?('-M') || ARGV.include?('-C'))
-  STDERR.print "Error: you must specify a command\n"
-  STDERR.print "(try '#{@app_name} --help')\n\n"
-  exit 1
-end
 
-# Load Rails app environment
-begin
-  require "config/environment"
-rescue LoadError
-  print_error "app environment"
-  raise
+if options.output
+  old_stdout = $stdout
+  begin
+     $stdout = File.open(options.output, "w")
+  rescue
+     $stderr.print "Error: Cannot write to #{options.output}\n\n"
+     exit 2
+  end
 end
 
-if ARGV.include?('-M')
-  generate_models
-elsif ARGV.include?('-C')
-  generate_controllers
+diagram.generate
+
+if options.out
+  $stdout = old_stdout
 end