lib/railroad in railroad-0.3.1 vs lib/railroad in railroad-0.3.2

@@ -14,19 +14,21 @@
 # (at your option) any later version.
 #
 
 APP_NAME  = "railroad"
 APP_HUMAN_NAME = "RailRoad"
-APP_VERSION   = [0,3,1]
+APP_VERSION   = [0,3,2]
 COPYRIGHT = "Copyright (C) 2007 Javier Smaldone"
 
 
+# Command line options parser
 class OptionsParser
 
   require 'optparse'
   require 'ostruct'
 
+  # Parse arguments from command line
   def self.parse(args)
 
     options = OpenStruct.new
 
     options.all = false
@@ -99,19 +101,19 @@
       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"
+           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"
+           STDERR.print "Error: Can't generate models AND controllers diagram\n\n"
            exit 1
          else 
            options.command = 'controllers'        
          end
       end
@@ -137,17 +139,18 @@
   end  # self.parse
  
   private 
 
   def self.option_error(msg, opts)
-    $stderr.print "Error: #{msg}\n\n #{opts}\n"
+    STDERR.print "Error: #{msg}\n\n #{opts}\n"
     exit 1
   end
 
 end  # class OptionsParser
 
 
+# Root class for RailRoad diagrams
 class AppDiagram
 
   def initialize(options)
     @options = options
     load_environment
@@ -158,37 +161,54 @@
   # Quotes a class name
   def node(name)
     '"' + name + '"'
   end
 
+  # Prevents Rails application from writing to STDOUT
+  def disable_stdout
+    @old_stdout = STDOUT.dup
+    STDOUT.reopen( PLATFORM =~ /mswin/ ? "NUL" : "/dev/null" )
+  end
+
+  # Restore STDOUT  
+  def enable_stdout
+    STDOUT.reopen(@old_stdout)
+  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
 
+  # Print error when loading application classes
   def print_error(type)
     STDERR.print "Error loading #{type}.\n  (Are you running " +
                  "#{APP_NAME} on the aplication's root directory?)\n\n"
   end
 
   # Load Rails application's environment
   def load_environment
     begin
+      disable_stdout
       require "config/environment"
+      enable_stdout
     rescue LoadError
+      enable_stdout
       print_error "app environment"
       raise
     end
   end
 
 end # class AppDiagram
 
 
+# RailRoad models diagram
 class ModelsDiagram < AppDiagram
 
   # Generate models diagram
   def generate
 
@@ -216,12 +236,15 @@
   private
 
   # Load model classes
   def load_classes
     begin
+      disable_stdout
       Dir.glob("app/models/**/*.rb") {|m| require m }
+      enable_stdout
     rescue LoadError
+      enable_stdout
       print_error "model classes"
       raise
     end
   end  # load_classes
 
@@ -314,10 +337,11 @@
   end # process_association
 
 end # class ModelsDiagram
 
 
+# RailRoad controllers diagram
 class ControllersDiagram < AppDiagram
 
   # Generate controllers diagram
   def generate
 
@@ -351,23 +375,27 @@
   private
 
   # Load controller classes
   def load_classes
     begin
+      disable_stdout
       # ApplicationController must be loaded first
       require "app/controllers/application.rb" 
       Dir.glob("app/controllers/**/*_controller.rb") do |c| 
         require c
       end
+      enable_stdout
     rescue LoadError
+      enable_stdout
       print_error "controller classes"
       raise
     end
   end # load_classes
 
   # Proccess controller class
   def process_class(current_class)
+
     # Print the node
     if @options.brief
       print "\t#{node(current_class.name)}\n"
     elsif current_class.is_a? Class 
       print "\t#{node(current_class.name)} " +
@@ -409,25 +437,25 @@
 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"
+  STDERR.print "Error: You must supply a command\n" +
+               "  (try #{APP_NAME} -h)\n\n"
   exit 1
 end
 
 if options.output
-  old_stdout = $stdout
+  old_stdout = STDOUT.dup
   begin
-     $stdout = File.open(options.output, "w")
+     STDOUT.reopen(options.output)
   rescue
-     $stderr.print "Error: Cannot write to #{options.output}\n\n"
+     STDERR.print "Error: Cannot write to #{options.output}\n\n"
      exit 2
   end
 end
 
 diagram.generate
 
-if options.out
-  $stdout = old_stdout
+if options.output
+  STDOUT.reopen(old_stdout)
 end