Fix ActiveRecord STI format lookup and guard __use__ cycles

lib/comma/extractor.rb

@@ -1,12 +1,15 @@
 # frozen_string_literal: true
 
 module Comma
+  class CircularStyleReference < StandardError; end
+
   class Extractor
     def initialize(instance, style, formats)
       @instance = instance
       @style = style
       @formats = formats
       @results = []
+      @style_stack = [style]
     end
 
     def results
@@ -19,8 +22,19 @@ def id(*args, &block)
     end
 
     def __use__(style)
-      # TODO: prevent infinite recursion
-      instance_eval(&@formats[style])
+      if @style_stack.include?(style)
+        chain = (@style_stack + [style]).join(' -> ')
+        raise Comma::CircularStyleReference, "Circular __use__ reference detected: #{chain}"
+      end
+
+      format = @formats.fetch(style) { raise "No comma format defined for style #{style}" }
+
+      @style_stack.push(style)
+      begin
+        instance_eval(&format)
+      ensure
+        @style_stack.pop
+      end
     end
 
     private

lib/comma/object.rb

@@ -4,19 +4,43 @@
 require 'comma/header_extractor'
 
 class Object
-  class_attribute :comma_formats
-
   class << self
     def comma(style = :default, &block)
-      (self.comma_formats ||= {})[style] = block
+      own_comma_formats[style] = block
+    end
+
+    def comma_formats
+      classes_with_own_formats.reverse_each.each_with_object({}) do |klass, formats|
+        formats.merge!(klass.instance_variable_get(:@own_comma_formats))
+      end
+    end
+
+    private
+
+    def own_comma_formats
+      @own_comma_formats ||= {}
     end
 
-    def inherited(subclass)
-      super
-      subclass.comma_formats = self.comma_formats ? self.comma_formats.dup : {}
+    # `self` here is always a Class: `comma`/`comma_formats` live on Object's
+    # singleton class, which only plain Class objects mirror, so a bare
+    # Module can never reach this method to begin with. `Class#superclass`
+    # in turn only ever yields another Class or nil, so `klass` stays a
+    # Class for the whole walk.
+    def classes_with_own_formats
+      classes = []
+      klass = self
+      while klass
+        classes << klass if klass.instance_variable_defined?(:@own_comma_formats)
+        klass = klass.superclass
+      end
+      classes
     end
   end
 
+  def comma_formats
+    self.class.comma_formats
+  end
+
   def to_comma(style = :default)
     extract_with(Comma::DataExtractor, style)
   end
@@ -28,13 +52,15 @@ def to_comma_headers(style = :default)
   private
 
   def extract_with(extractor_class, style = :default)
-    raise_unless_style_exists(style)
-    extractor_class.new(self, style, self.comma_formats).results
+    formats = comma_formats
+    raise_unless_style_exists(formats, style)
+    extractor_class.new(self, style, formats).results
   end
 
-  def raise_unless_style_exists(style)
-    return if self.comma_formats && self.comma_formats[style]
+  def raise_unless_style_exists(formats, style)
+    return if formats[style]
 
-    raise "No comma format for class #{self.class} defined for style #{style}"
+    receiver = is_a?(Module) ? self : self.class
+    raise "No comma format for class #{receiver} defined for style #{style}"
   end
 end

spec/comma/comma_spec.rb

@@ -263,7 +263,7 @@ def initialize(content, created_at = Time.now, updated_at = Time.now)
     end
   end
 
-  describe 'on objects using Single Table Inheritance' do
+  describe 'on objects using Single Table Inheritance' do # rubocop:disable Metrics/BlockLength
     before do
       class MySuperClass
         attr_accessor :content
@@ -296,6 +296,29 @@ class ChildClassNoComma < MySuperClass
     it 'should return and array of data content, as defined in comma block in super class, if not present in child' do
       expect(@childNoComma.to_comma).to eq(%w[super-content])
     end
+
+    it 'should reflect changes to the superclass format made after the subclass was defined' do
+      class ReopenedSuperClass
+        attr_accessor :content
+        comma do; content end
+
+        def initialize(content)
+          @content = 'super-' + content
+        end
+      end
+
+      class ReopenedChildNoComma < ReopenedSuperClass
+      end
+
+      ReopenedSuperClass.class_eval do
+        comma do
+          content(&:upcase)
+        end
+      end
+
+      child = ReopenedChildNoComma.new('content')
+      expect(child.to_comma).to eq(%w[SUPER-CONTENT])
+    end
   end
 end
 
@@ -322,3 +345,30 @@ class ChildClassNoComma < MySuperClass
   its(:size) { should eq(3) }
   it { should eq(['Programming Ruby', 'Foo, Inc.', 'The Pickaxe book']) }
 end
+
+describe Comma, '__use__ keyword with a circular reference' do
+  it 'should raise Comma::CircularStyleReference instead of overflowing the stack' do
+    obj = Class.new(Struct.new(:id, :title)) do
+      comma :a do
+        title
+        __use__ :b
+      end
+
+      comma :b do
+        __use__ :a
+      end
+    end.new(1, 'Programming Ruby')
+
+    expect { obj.to_comma(:a) }.to raise_error(Comma::CircularStyleReference, /a -> b -> a/)
+  end
+
+  it 'should raise Comma::CircularStyleReference for direct self-reference' do
+    obj = Class.new(Struct.new(:id)) do
+      comma :a do
+        __use__ :a
+      end
+    end.new(1)
+
+    expect { obj.to_comma(:a) }.to raise_error(Comma::CircularStyleReference, /a -> a/)
+  end
+end

spec/comma/rails/active_record_spec.rb

@@ -191,7 +191,6 @@ class Cat < Animal
       expect(@dog.to_comma).to eq %w[Dog-Rex]
     end
 
-    # FIXME: this one is failing - the comma block from Dog is executed instead of the one from the super class
     it 'should return and array of data content, as defined in comma block in super class, if not present in child' do
       expect(@cat.to_comma).to eq %w[Super-Kitty]
     end