lib/rails_erd/domain.rb in rails-erd-0.2.0 vs lib/rails_erd/domain.rb in rails-erd-0.3.0

@@ -1,15 +1,22 @@
 require "set"
 require "rails_erd"
 require "rails_erd/entity"
 require "rails_erd/relationship"
-require "rails_erd/relationship/cardinality"
 require "rails_erd/attribute"
 
 module RailsERD
   # The domain describes your Rails domain model. This class is the starting
   # point to get information about your models.
+  #
+  # === Options
+  #
+  # The following options are available:
+  #
+  # warn:: When set to +false+, no warnings are printed to the
+  #        command line while processing the domain model. Defaults
+  #        to +true+.
   class Domain
     class << self
       # Generates a domain model object based on all loaded subclasses of
       # <tt>ActiveRecord::Base</tt>. Make sure your models are loaded before calling
       # this method.
@@ -36,21 +43,21 @@
       defined? Rails and Rails.application and Rails.application.class.parent.name
     end
     
     # Returns all entities of your domain model.
     def entities
-      @entities ||= entity_mapping.values.sort
+      @entities ||= Entity.from_models(self, @models)
     end
     
     # Returns all relationships in your domain model.
     def relationships
       @relationships ||= Relationship.from_associations(self, associations)
     end
     
     # Returns a specific entity object for the given Active Record model.
     def entity_for(model) # @private :nodoc:
-      entity_mapping[model] or raise "model #{model} exists, but is not included in the domain"
+      entity_mapping[model] or raise "model #{model} exists, but is not included in domain"
     end
     
     # Returns an array of relationships for the given Active Record model.
     def relationships_for(model) # @private :nodoc:
       relationships_mapping[model] or []
@@ -59,14 +66,22 @@
     def inspect # @private :nodoc:
       "#<#{self.class}:0x%.14x {%s}>" %
         [object_id << 1, relationships.map { |rel| "#{rel.source} => #{rel.destination}" } * ", "]
     end
     
+    def warn(message) # @private :nodoc:
+      puts "Warning: #{message}" if options.warn
+    end
+    
     private
     
     def entity_mapping
-      @entity_mapping ||= Hash[@models.collect { |model| [model, Entity.new(self, model)] }]
+      @entity_mapping ||= {}.tap do |mapping|
+        entities.each do |entity|
+          mapping[entity.model] = entity
+        end
+      end
     end
     
     def relationships_mapping
       @relationships_mapping ||= {}.tap do |mapping|
         relationships.each do |relationship|
@@ -88,14 +103,10 @@
       model = association.klass
       
       # Raises error if model is not in the domain.
       entity_for model
     rescue => e
-      warn "Invalid association #{association_description(association)} (#{e.message})"
-    end
-    
-    def warn(message)
-      puts "Warning: #{message}" unless options.suppress_warnings
+      warn "Ignoring invalid association #{association_description(association)} (#{e.message})"
     end
     
     def association_description(association)
       "#{association.name.inspect} on #{association.active_record}"
     end