Define the methods in a module and prepend it so that when a model overrides a method, we can still use the super behaviour

Author: Kyle d'Oliveira

lib/polymorphic_integer_type/extensions.rb

@@ -24,26 +24,32 @@ def belongs_to(name, scope = nil, **options)
             mapping
           end
 
-          define_method foreign_type do
-            t = super()
-            mapping[t]
-          end
+          foreign_type_extension = Module.new do
+            define_method foreign_type do
+              t = super()
+              self.class.send("#{foreign_type}_mapping")[t]
+            end
 
-          define_method "#{foreign_type}=" do |klass|
-            enum = mapping.key(klass.to_s)
-            if klass.kind_of?(Class) && klass <= ActiveRecord::Base
-              enum ||= mapping.key(klass.base_class.to_s)
-              enum ||= mapping.key(klass.base_class.sti_name)
+            define_method "#{foreign_type}=" do |klass|
+              mapping = self.class.send("#{foreign_type}_mapping")
+              enum = mapping.key(klass.to_s)
+              if klass.kind_of?(Class) && klass <= ActiveRecord::Base
+                enum ||= mapping.key(klass.sti_name)
+                enum ||= mapping.key(klass.base_class.to_s)
+                enum ||= mapping.key(klass.base_class.sti_name)
+              end
+              enum ||= klass if klass != NilClass
+              super(enum)
             end
-            enum ||= klass if klass != NilClass
-            super(enum)
-          end
 
-          define_method "#{name}=" do |record|
-            super(record)
-            send("#{foreign_type}=", record.class)
+            define_method "#{name}=" do |record|
+              super(record)
+              send("#{foreign_type}=", record.class)
+            end
           end
 
+          include(foreign_type_extension)
+
           validate do
             t = send(foreign_type)
             unless t.nil? || mapping.values.include?(t)

spec/support/link.rb

@@ -3,4 +3,12 @@ class Link < ActiveRecord::Base
 
   belongs_to :source, polymorphic: true, integer_type: true
   belongs_to :target, polymorphic: true, integer_type: true
+
+  def source=(val)
+    super(val)
+  end
+
+  def source_type=(val)
+    super(val)
+  end
 end