option to group STI column annotation

README.md

@@ -251,6 +251,7 @@ you can do so with a simple environment variable, instead of editing the
             --ignore-unknown-models      don't display warnings for bad model files
             --with-comment               include database comments in model annotations
             --with-comment-column        include database comments in model annotations, as its own column, after all others
+            --group-sti-columns          group annotation columns by STI class ownership
 
 ### Option: `additional_file_patterns`
 

lib/annotate/annotate_models.rb

@@ -3,6 +3,7 @@
 require 'bigdecimal'
 
 require 'annotate/constants'
+require 'annotate/sti_columns'
 require_relative 'annotate_models/file_patterns'
 
 module AnnotateModels
@@ -171,27 +172,43 @@ def get_schema_info(klass, header, options = {}) # rubocop:disable Metrics/Metho
       # Output annotation
       bare_max_attrs_length = cols_meta.map { |_, m| m[:simple_formatted_attrs].length }.max
 
-      cols.each do |col|
-        col_type = cols_meta[col.name][:col_type]
-        attrs = cols_meta[col.name][:attrs]
-        col_name = cols_meta[col.name][:col_name]
-        simple_formatted_attrs = cols_meta[col.name][:simple_formatted_attrs]
-        col_comment = cols_meta[col.name][:col_comment]
-
-        if options[:format_rdoc]
-          info << sprintf("# %-#{max_size}.#{max_size}s<tt>%s</tt>", "*#{col_name}*::", attrs.unshift(col_type).join(", ")).rstrip + "\n"
-        elsif options[:format_yard]
-          info << sprintf("# @!attribute #{col_name}") + "\n"
-          ruby_class = col.respond_to?(:array) && col.array ? "Array<#{map_col_type_to_ruby_classes(col_type)}>": map_col_type_to_ruby_classes(col_type)
-          info << sprintf("#   @return [#{ruby_class}]") + "\n"
-        elsif options[:format_markdown]
-          name_remainder = max_size - col_name.length - non_ascii_length(col_name)
-          type_remainder = (md_type_allowance - 2) - col_type.length
-          info << (sprintf("# **`%s`**%#{name_remainder}s | `%s`%#{type_remainder}s | `%s`", col_name, " ", col_type, " ", attrs.join(", ").rstrip)).gsub('``', '  ').rstrip + "\n"
-        elsif with_comments_column
-          info << format_default(col_name, max_size, col_type, bare_type_allowance, simple_formatted_attrs, bare_max_attrs_length, col_comment)
-        else
-          info << format_default(col_name, max_size, col_type, bare_type_allowance, simple_formatted_attrs)
+      if options[:group_sti_columns]
+        grouped_cols = Annotate::StiColumns.partition(klass, cols)
+      else
+        grouped_cols = [[nil, cols]]
+      end
+
+      grouped_cols.each do |group_label, group_columns|
+        if group_label
+          if options[:format_markdown]
+            info << "#\n# ### #{group_label}\n#\n"
+          else
+            info << "#\n# -- #{group_label} --\n#\n"
+          end
+        end
+
+        group_columns.each do |col|
+          col_type = cols_meta[col.name][:col_type]
+          attrs = cols_meta[col.name][:attrs]
+          col_name = cols_meta[col.name][:col_name]
+          simple_formatted_attrs = cols_meta[col.name][:simple_formatted_attrs]
+          col_comment = cols_meta[col.name][:col_comment]
+
+          if options[:format_rdoc]
+            info << sprintf("# %-#{max_size}.#{max_size}s<tt>%s</tt>", "*#{col_name}*::", attrs.unshift(col_type).join(", ")).rstrip + "\n"
+          elsif options[:format_yard]
+            info << sprintf("# @!attribute #{col_name}") + "\n"
+            ruby_class = col.respond_to?(:array) && col.array ? "Array<#{map_col_type_to_ruby_classes(col_type)}>": map_col_type_to_ruby_classes(col_type)
+            info << sprintf("#   @return [#{ruby_class}]") + "\n"
+          elsif options[:format_markdown]
+            name_remainder = max_size - col_name.length - non_ascii_length(col_name)
+            type_remainder = (md_type_allowance - 2) - col_type.length
+            info << (sprintf("# **`%s`**%#{name_remainder}s | `%s`%#{type_remainder}s | `%s`", col_name, " ", col_type, " ", attrs.join(", ").rstrip)).gsub('``', '  ').rstrip + "\n"
+          elsif with_comments_column
+            info << format_default(col_name, max_size, col_type, bare_type_allowance, simple_formatted_attrs, bare_max_attrs_length, col_comment)
+          else
+            info << format_default(col_name, max_size, col_type, bare_type_allowance, simple_formatted_attrs)
+          end
         end
       end
 
@@ -733,7 +750,7 @@ def annotate_model_file(annotated, file, header, options)
         klass = get_model_class(file)
         do_annotate = klass.is_a?(Class) &&
           klass < ActiveRecord::Base &&
-          (!options[:exclude_sti_subclasses] || !(klass.superclass < ActiveRecord::Base && klass.table_name == klass.superclass.table_name)) &&
+          (!options[:exclude_sti_subclasses] || !Annotate::StiColumns.subclass?(klass)) &&
           !klass.abstract_class? &&
           klass.table_exists?
 

lib/annotate/constants.rb

@@ -18,7 +18,7 @@ module Constants
       :trace, :timestamp, :exclude_serializers, :classified_sort,
       :show_foreign_keys, :show_complete_foreign_keys,
       :exclude_scaffolds, :exclude_controllers, :exclude_helpers,
-      :exclude_sti_subclasses, :ignore_unknown_models, :with_comment, :with_comment_column,
+      :exclude_sti_subclasses, :group_sti_columns, :ignore_unknown_models, :with_comment, :with_comment_column,
       :show_check_constraints
     ].freeze
 

lib/annotate/parser.rb

@@ -309,6 +309,11 @@ def add_options_to_parser(option_parser) # rubocop:disable Metrics/MethodLength,
                        "include database comments in model annotations, as its own column, after all others") do
         env['with_comment_column'] = 'true'
       end
+
+      option_parser.on('--group-sti-columns',
+                       "group annotation columns by STI class ownership") do
+        env['group_sti_columns'] = 'true'
+      end
     end
   end
 end

lib/annotate/sti_columns.rb

@@ -0,0 +1,121 @@
+require 'set'
+
+module Annotate
+  module StiColumns
+    class << self
+      def subclass?(klass)
+        klass.superclass < ActiveRecord::Base &&
+          klass.table_name == klass.superclass.table_name
+      end
+
+      def base_class?(klass)
+        klass.column_names.include?(klass.inheritance_column) &&
+          sti_descendants(klass).any?
+      end
+
+      def columns_referenced_in(klass)
+        cols = Set.new
+        cols.merge(klass.validators.flat_map { |v| v.attributes.map(&:to_s) })
+        if klass.respond_to?(:reflect_on_all_associations)
+          cols.merge(
+            klass.reflect_on_all_associations(:belongs_to).flat_map { |a|
+              fk = [a.foreign_key.to_s]
+              fk << "#{a.name}_type" if a.respond_to?(:polymorphic?) && a.polymorphic?
+              fk
+            }
+          )
+        end
+        cols.merge(klass.defined_enums.keys) if klass.respond_to?(:defined_enums)
+        if klass.respond_to?(:stored_attributes) && klass.stored_attributes.any?
+          cols.merge(klass.stored_attributes.keys.map(&:to_s))
+        end
+        cols
+      end
+
+      def columns_owned_by(klass)
+        return Set.new unless subclass?(klass)
+
+        columns_referenced_in(klass) - columns_referenced_in(klass.superclass)
+      end
+
+      def partition(klass, cols)
+        if subclass?(klass)
+          partition_for_subclass(klass, cols)
+        elsif base_class?(klass)
+          partition_for_base_class(klass, cols)
+        else
+          [[nil, cols]]
+        end
+      end
+
+      private
+
+      # Returns all descendants of klass that share the same table (STI).
+      # Uses descendants (all levels) rather than subclasses (direct only)
+      # to support multi-level STI hierarchies.
+      #
+      # Attempts eager loading first to ensure all subclasses are visible.
+      def sti_descendants(klass)
+        ensure_models_loaded
+        if klass.respond_to?(:descendants)
+          klass.descendants.select { |d| d.table_name == klass.table_name }
+        else
+          klass.subclasses.select { |d| d.table_name == klass.table_name }
+        end
+      end
+
+      def ensure_models_loaded
+        return if @models_loaded
+
+        if defined?(Rails) && Rails.respond_to?(:application) && Rails.application
+          Rails.application.eager_load! unless Rails.application.config.eager_load
+        end
+        @models_loaded = true
+      end
+
+      def partition_for_subclass(klass, cols)
+        owned = columns_owned_by(klass)
+        base_name = klass.superclass.name.demodulize
+        shared = cols.reject { |c| owned.include?(c.name.to_s) }
+        specific = cols.select { |c| owned.include?(c.name.to_s) }
+
+        if specific.empty?
+          $stderr.puts "Warning: --group-sti-columns could not determine which columns belong to #{klass.name}."
+          $stderr.puts "  Add validations, associations, or enums to #{klass.name} to improve grouping."
+        end
+
+        groups = []
+        groups << ["#{base_name} columns", shared] if shared.any?
+        groups << ["#{klass.name.demodulize} columns", specific] if specific.any?
+        groups
+      end
+
+      def partition_for_base_class(klass, cols)
+        all_descendants = sti_descendants(klass).sort_by(&:name)
+        ownership = {}
+
+        all_descendants.each do |sub|
+          columns_owned_by(sub).each do |col_name|
+            ownership[col_name] ||= sub.name.demodulize
+          end
+        end
+
+        base_cols = cols.reject { |c| ownership.key?(c.name.to_s) }
+        groups = [["#{klass.name.demodulize} columns", base_cols]]
+
+        all_descendants.each do |sub|
+          sub_name = sub.name.demodulize
+          sub_cols = cols.select { |c| ownership[c.name.to_s] == sub_name }
+          groups << ["#{sub_name} columns", sub_cols] if sub_cols.any?
+        end
+
+        if ownership.empty?
+          $stderr.puts "Warning: --group-sti-columns found STI subclasses for #{klass.name} but no columns could be assigned."
+          $stderr.puts "  Add validations, associations, or enums to subclasses to improve grouping."
+        end
+
+        groups
+      end
+    end
+  end
+end

lib/generators/annotate/templates/auto_annotate_models.rake

@@ -34,6 +34,7 @@ if Rails.env.development?
       'exclude_controllers'         => 'true',
       'exclude_helpers'             => 'true',
       'exclude_sti_subclasses'      => 'false',
+      'group_sti_columns'           => 'false',
       'ignore_model_sub_dir'        => 'false',
       'ignore_columns'              => nil,
       'ignore_routes'               => nil,

lib/tasks/annotate_models.rake

@@ -35,6 +35,7 @@ task annotate_models: :environment do
   options[:exclude_controllers] = Annotate::Helpers.true?(ENV.fetch('exclude_controllers', 'true'))
   options[:exclude_helpers] = Annotate::Helpers.true?(ENV.fetch('exclude_helpers', 'true'))
   options[:exclude_sti_subclasses] = Annotate::Helpers.true?(ENV['exclude_sti_subclasses'])
+  options[:group_sti_columns] = Annotate::Helpers.true?(ENV['group_sti_columns'])
   options[:ignore_model_sub_dir] = Annotate::Helpers.true?(ENV['ignore_model_sub_dir'])
   options[:format_bare] = Annotate::Helpers.true?(ENV['format_bare'])
   options[:format_rdoc] = Annotate::Helpers.true?(ENV['format_rdoc'])