Extend override check to also check arity

Override now cares about parents method arity: if parents arity
doesn't match raises an error.

app/models/dashboard_group_milestone.rb

@@ -5,7 +5,6 @@ class DashboardGroupMilestone < GlobalMilestone
 
   attr_reader :group_name
 
-  override :initialize
   def initialize(milestone)
     super(milestone.title, Array(milestone))
 

changelogs/unreleased/42125-extend-override-check-to-also-check-arity.yml

+---
+title: Extend override check to also check arity
+merge_request: 23498
+author: Jacopo Beschi @jacopo-beschi
+type: added

lib/gitlab/utils/override.rb

@@ -4,16 +4,11 @@ module Gitlab
   module Utils
     module Override
       class Extension
-        def self.verify_class!(klass, method_name)
-          instance_method_defined?(klass, method_name) ||
-            raise(
-              NotImplementedError.new(
-                "#{klass}\##{method_name} doesn't exist!"))
-        end
-
-        def self.instance_method_defined?(klass, name, include_super: true)
-          klass.instance_methods(include_super).include?(name) ||
-            klass.private_instance_methods(include_super).include?(name)
+        def self.verify_class!(klass, method_name, arity)
+          extension = new(klass)
+          parents = extension.parents_for(klass)
+          extension.verify_method!(
+            klass: klass, parents: parents, method_name: method_name, sub_method_arity: arity)
         end
 
         attr_reader :subject
@@ -22,35 +17,77 @@ module Gitlab
           @subject = subject
         end
 
-        def add_method_name(method_name)
-          method_names << method_name
-        end
-
-        def add_class(klass)
-          classes << klass
+        def parents_for(klass)
+          index = klass.ancestors.index(subject)
+          klass.ancestors.drop(index + 1)
         end
 
         def verify!
           classes.each do |klass|
-            index = klass.ancestors.index(subject)
-            parents = klass.ancestors.drop(index + 1)
+            parents = parents_for(klass)
+
+            method_names.each_pair do |method_name, arity|
+              verify_method!(
+                klass: klass,
+                parents: parents,
+                method_name: method_name,
+                sub_method_arity: arity)
+            end
+          end
+        end
+
+        def verify_method!(klass:, parents:, method_name:, sub_method_arity:)
+          overridden_parent = parents.find do |parent|
+            instance_method_defined?(parent, method_name)
+          end
+
+          raise NotImplementedError.new("#{klass}\##{method_name} doesn't exist!") unless overridden_parent
 
-            method_names.each do |method_name|
-              parents.any? do |parent|
-                self.class.instance_method_defined?(
-                  parent, method_name, include_super: false)
-              end ||
-                raise(
-                  NotImplementedError.new(
-                    "#{klass}\##{method_name} doesn't exist!"))
+          super_method_arity = find_direct_method(overridden_parent, method_name).arity
+
+          unless arity_compatible?(sub_method_arity, super_method_arity)
+            raise NotImplementedError.new("#{subject}\##{method_name} has arity of #{sub_method_arity}, but #{overridden_parent}\##{method_name} has arity of #{super_method_arity}")
+          end
         end
+
+        def add_method_name(method_name, arity = nil)
+          method_names[method_name] = arity
+        end
+
+        def add_class(klass)
+          classes << klass
         end
+
+        def verify_override?(method_name)
+          method_names.has_key?(method_name)
         end
 
         private
 
+        def instance_method_defined?(klass, name)
+          klass.instance_methods(false).include?(name) ||
+            klass.private_instance_methods(false).include?(name)
+        end
+
+        def find_direct_method(klass, name)
+          method = klass.instance_method(name)
+          method = method.super_method until method && klass == method.owner
+          method
+        end
+
+        def arity_compatible?(sub_method_arity, super_method_arity)
+          if sub_method_arity >= 0 && super_method_arity >= 0
+            # Regular arguments
+            sub_method_arity == super_method_arity
+          else
+            # It's too complex to check this case, just allow sub-method having negative arity
+            # But we don't allow sub_method_arity > 0 yet super_method_arity < 0
+            sub_method_arity < 0
+          end
+        end
+
         def method_names
-          @method_names ||= []
+          @method_names ||= {}
         end
 
         def classes
@@ -80,12 +117,22 @@ module Gitlab
       def override(method_name)
         return unless ENV['STATIC_VERIFICATION']
 
-        if is_a?(Class)
-          Extension.verify_class!(self, method_name)
-        else # We delay the check for modules
         Override.extensions[self] ||= Extension.new(self)
         Override.extensions[self].add_method_name(method_name)
       end
+
+      def method_added(method_name)
+        super
+
+        return unless ENV['STATIC_VERIFICATION']
+        return unless Override.extensions[self]&.verify_override?(method_name)
+
+        method_arity = instance_method(method_name).arity
+        if is_a?(Class)
+          Extension.verify_class!(self, method_name, method_arity)
+        else # We delay the check for modules
+          Override.extensions[self].add_method_name(method_name, method_arity)
+        end
       end
 
       def included(base = nil)

spec/lib/gitlab/utils/override_spec.rb

@@ -25,13 +25,23 @@ describe Gitlab::Utils::Override do
 
   let(:klass) { subject }
 
-  def good(mod)
+  def good(mod, bad_arity: false, negative_arity: false)
     mod.module_eval do
       override :good
+
+      if bad_arity
+        def good(num)
+        end
+      elsif negative_arity
+        def good(*args)
+          super.succ
+        end
+      else
         def good
           super.succ
         end
       end
+    end
 
     mod
   end
@@ -56,6 +66,14 @@ describe Gitlab::Utils::Override do
       described_class.verify!
     end
 
+    it 'checks ok for overriding method using negative arity' do
+      good(subject, negative_arity: true)
+      result = instance.good
+
+      expect(result).to eq(1)
+      described_class.verify!
+    end
+
     it 'raises NotImplementedError when it is not overriding anything' do
       expect do
         bad(subject)
@@ -63,6 +81,14 @@ describe Gitlab::Utils::Override do
         described_class.verify!
       end.to raise_error(NotImplementedError)
     end
+
+    it 'raises NotImplementedError when overriding a method with different arity' do
+      expect do
+        good(subject, bad_arity: true)
+        instance.good(1)
+        described_class.verify!
+      end.to raise_error(NotImplementedError)
+    end
   end
 
   shared_examples 'checking as intended, nothing was overridden' do