Pass index to enhancement's constructor

Shopify · Oct 22, 2024 · ba3e595 · ba3e595
1 parent 1b2b113
commit ba3e595
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/lib/ruby_indexer/lib/ruby_indexer/enhancement.rb b/lib/ruby_indexer/lib/ruby_indexer/enhancement.rb
@@ -8,12 +8,16 @@ class Enhancement
 
     abstract!
 
+    sig { params(index: Index).void }
+    def initialize(index)
+      @index = index
+    end
+
     # The `on_extend` indexing enhancement is invoked whenever an extend is encountered in the code. It can be used to
     # register for an included callback, similar to what `ActiveSupport::Concern` does in order to auto-extend the
     # `ClassMethods` modules
     sig do
       overridable.params(
-        index: Index,
         owner: T.nilable(Entry::Namespace),
         node: Prism::CallNode,
         file_path: String,
@@ -23,11 +27,10 @@ class Enhancement
         ),
       ).void
     end
-    def on_call_node_enter(index, owner, node, file_path, code_units_cache); end
+    def on_call_node_enter(owner, node, file_path, code_units_cache); end
 
     sig do
       overridable.params(
-        index: Index,
         owner: T.nilable(Entry::Namespace),
         node: Prism::CallNode,
         file_path: String,
@@ -37,6 +40,6 @@ def on_call_node_enter(index, owner, node, file_path, code_units_cache); end
         ),
       ).void
     end
-    def on_call_node_leave(index, owner, node, file_path, code_units_cache); end
+    def on_call_node_leave(owner, node, file_path, code_units_cache); end
   end
 end
diff --git a/lib/ruby_indexer/test/enhancements_test.rb b/lib/ruby_indexer/test/enhancements_test.rb
@@ -48,7 +48,7 @@ def on_call_node_enter(index, owner, node, file_path, code_units_cache)
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
       index(<<~RUBY)
         module ActiveSupport
           module Concern
@@ -124,7 +124,7 @@ def on_call_node_enter(index, owner, node, file_path, code_units_cache)
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
       index(<<~RUBY)
         module ActiveSupport
           module Concern
@@ -169,7 +169,7 @@ def name
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
 
       _stdout, stderr = capture_io do
         index(<<~RUBY)
@@ -204,7 +204,7 @@ def name
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
 
       _stdout, stderr = capture_io do
         index(<<~RUBY)