lib/modaldiagrams/modaldiagrams.rb in modaldiagrams-1.1.5 vs lib/modaldiagrams/modaldiagrams.rb in modaldiagrams-1.1.6

@@ -33,14 +33,16 @@
       relation_classes = [] # each element is the 2-element array of class names of the corresponding relation in relations
 
       model_selection_options = {
           :all_models => cfg.include_all_models,
           :dynamic_models => cfg.include_dynamic_models,
-          :include_files => false
+          :include_files => false,
+          :exclude_models => cfg.exclude_models
       }
 
-      models = dbmodels(model_selection_options.merge(:exclude_sti_models => true))
+      models, excluded_models = dbmodels(model_selection_options.merge(:exclude_sti_models => true))
+      excluded_class_names = excluded_models.map(&:name)
 
       models.each do |cls|
         if cls.respond_to?(:reflect_on_all_associations) && ActiveRecord::Base.connection.table_exists?(cls.table_name)
           # Note: Don't use content_columns ignores columns ending with _id which I use for enum fields
           columns = cls.columns.reject { |c| c.primary || c.name =~ /(_count)$/ || c.name == cls.inheritance_column || c.name =~ /^(created_at|updated_at)$/ }.map{|c| field_spec(cfg, c)}
@@ -70,10 +72,11 @@
           end
           classes[cluster] << %{"#{cls}" [shape=Mrecord, label="{#{cls}|#{columns.join('\\l')}\\l}"]}
           cluster_classes[cluster] ||= []
           cluster_classes[cluster] << cls.to_s
           cls.reflect_on_all_associations.each do |assoc|
+            next if assoc.class_name.in?(excluded_class_names)
             target,type = nil,nil
             case assoc.macro
             when :has_many
               unless assoc.options[:through]
                 target = assoc.class_name
@@ -137,11 +140,11 @@
           end
         end
       end
 
       if cfg.show_sti
-        sti_classes = dbmodels(model_selection_options.merge(:exclude_non_sti_models => true))
+        sti_classes, exc = dbmodels(model_selection_options.merge(:exclude_non_sti_models => true))
         sti_classes.each do |sti_class|
           cls = sti_class.base_class
           if cls.respond_to?(:cluster)
             cluster = cls.cluster.to_s
           else
@@ -165,11 +168,11 @@
           relation_classes << [base_class.to_s, sti_class.to_s]
         end
       end
 
       fn = Rails.root.join('db/diagrams/diagram.dot')
-      mkdir_p fn.dirname
+      FileUtils.mkdir_p fn.dirname
       File.open(fn,'w') do |f|
         add_diagram_header f
         cluster_id = 0
         all_classes = []
         classes.keys.each do |cluster|
@@ -203,10 +206,11 @@
     # Options:
     #   :all_models                # Return also models in plugins, not only in the app (app/models)
     #   :dynamic_models            # Return dynamically defined models too (not defined in a model file)
     #   :exclude_sti_models        # Exclude derived (STI) models
     #   :exclude_non_sti_models    # Exclude top level models
+    #   :exclude_models            # Array of models to exclude from the diagrams
     #   :include_files             # Return also the model definition file pathnames (return pairs of [model, file])
     #   :only_app_files            #   But return nil for files not in the app proper
     #   :only_app_tree_files       #   But return nil for files not in the app directory tree (app, vendor...)
     def dbmodels(options={})
 
@@ -245,12 +249,21 @@
       models = models.uniq.reject{|model| !has_table?(model)}
 
       non_sti_models, sti_models = models.partition{|model| model.base_class==model}
 
       models = []
-      models += non_sti_models unless options[:exclude_non_sti_models]
-      models += sti_models unless options[:exclude_sti_models]
+      excluded_models = []
+      if options[:exclude_non_sti_models]
+        excluded_models += non_sti_models
+      else
+        models += non_sti_models
+      end
+      if options[:exclude_sti_models]
+        excluded_models += sti_models
+      else
+        models += sti_models
+      end
       if options[:include_files]
         models = models.map{|model| [model, files[model.to_s]]}
         if options[:only_app_files] || options[:only_app_tree_files]
           if options[:only_app_files]
             suffix = models_dir.to_s
@@ -259,10 +272,15 @@
           end
           suffix += '/' unless suffix.ends_with?('/')
           models = models.map{|model, file| [model, file && (file.starts_with?(suffix) ? file : nil)]}
         end
       end
-      models
+      excluded_models = []
+      if options[:exclude_models].present?
+        excluded_models = Array(options[:exclude_models]).map{|m| String===m ? m.constantize : m}
+        models -= excluded_models
+      end
+      [models, excluded_models]
     end
 
     def has_table?(cls)
      (cls != ActiveRecord::Base) && cls.respond_to?(:table_name) && cls.table_name.present?
     end