Allow memoization when included modules prepend MemoWise and define initialize

Fixes #302

Co-authored-by: alpaca-tc <alpaca-tc@alpaca.tc>

Author: JacobEvelyn

CHANGELOG.md

@@ -9,6 +9,8 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 **Gem enhancements:** none
 
+- Allow memoization when `include`d modules `prepend MemoWise` and define `initialize` [[#327]](https://github.com/panorama-ed/memo_wise/pull/327)
+
 _No breaking changes!_
 
 **Project enhancements:** none

lib/memo_wise.rb

@@ -53,10 +53,9 @@ module MemoWise
   # [this article](https://www.ruby-lang.org/en/news/2019/12/12/separation-of-positional-and-keyword-arguments-in-ruby-3-0/)
   # for more information.
   #
+
   # :nocov:
-  all_args = RUBY_VERSION < "2.7" ? "*" : "..."
-  # :nocov:
-  class_eval <<~HEREDOC, __FILE__, __LINE__ + 1
+  INITIALIZE_LITERAL = <<~HEREDOC
     # On Ruby 2.7 or greater:
     #
     # def initialize(...)
@@ -71,11 +70,14 @@ module MemoWise
     #   super
     # end
 
-    def initialize(#{all_args})
+    def initialize(#{RUBY_VERSION < '2.7' ? '*' : '...'})
       MemoWise::InternalAPI.create_memo_wise_state!(self)
       super
     end
   HEREDOC
+  # :nocov:
+
+  class_eval(INITIALIZE_LITERAL, __FILE__, __LINE__ + 1)
 
   module CreateMemoWiseStateOnExtended
     def extended(base)
@@ -93,6 +95,13 @@ def inherited(subclass)
   end
   private_constant(:CreateMemoWiseStateOnInherited)
 
+  module CreateMemoWiseStateOnIncluded
+    def included(base)
+      base.prepend(MemoWise)
+    end
+  end
+  private_constant(:CreateMemoWiseStateOnIncluded)
+
   # @private
   #
   # Private setup method, called automatically by `prepend MemoWise` in a class.
@@ -154,10 +163,7 @@ def memo_wise(method_name_or_hash)
             klass.singleton_class.prepend(CreateMemoWiseStateOnExtended)
           end
         when Hash
-          unless method_name_or_hash.keys == [:self]
-            raise ArgumentError,
-                  "`:self` is the only key allowed in memo_wise"
-          end
+          raise ArgumentError, "`:self` is the only key allowed in memo_wise" unless method_name_or_hash.keys == [:self]
 
           method_name = method_name_or_hash[:self]
 
@@ -175,6 +181,8 @@ def memo_wise(method_name_or_hash)
         if klass.is_a?(Class) && !klass.singleton_class?
           klass.singleton_class.prepend(CreateMemoWiseStateOnInherited)
         else
+          klass.singleton_class.prepend(CreateMemoWiseStateOnIncluded) if klass.is_a?(Module) && !klass.singleton_class?
+
           klass.prepend(CreateMemoWiseStateOnInherited)
         end
 

spec/memo_wise_spec.rb

@@ -351,11 +351,32 @@ def module2_method
           end
         end
 
+        let(:klass_with_initializer) do
+          Class.new do
+            include Module1
+            def initialize(*); end
+          end
+        end
+
+        let(:module_with_initializer) do
+          Module.new do
+            include Module1
+            def initialize(*); end
+          end
+        end
+
+        let(:klass_with_module_with_initializer) do
+          Class.new do
+            include Module3
+          end
+        end
+
         let(:instance) { klass.new }
 
         before(:each) do
           stub_const("Module1", module1)
           stub_const("Module2", module2)
+          stub_const("Module3", module_with_initializer)
         end
 
         it "memoizes inherited methods separately" do
@@ -364,6 +385,32 @@ def module2_method
           expect(Array.new(4) { instance.module2_method }).to all eq("module2_method")
           expect(instance.module2_method_counter).to eq(1)
         end
+
+        it "can memoize klass with initializer" do
+          instance = klass_with_initializer.new(true)
+          expect { instance.module1_method }.not_to raise_error
+
+          expect(Array.new(4) { instance.module1_method }).to all eq("module1_method")
+          expect(instance.module1_method_counter).to eq(1)
+        end
+
+        it "can memoize klass with module with initializer" do
+          instance = klass_with_module_with_initializer.new(true)
+          expect { instance.module1_method }.not_to raise_error
+
+          expect(Array.new(4) { instance.module1_method }).to all eq("module1_method")
+          expect(instance.module1_method_counter).to eq(1)
+        end
+
+        it "can reset klass with initializer" do
+          instance = klass_with_initializer.new(true)
+          expect { instance.reset_memo_wise }.not_to raise_error
+        end
+
+        it "can reset klass with module with initializer" do
+          instance = klass_with_module_with_initializer.new(true)
+          expect { instance.reset_memo_wise }.not_to raise_error
+        end
       end
 
       context "when the class, its superclass, and its module all memoize methods" do

spec/proxying_original_method_params_spec.rb

@@ -1,27 +1,10 @@
 # frozen_string_literal: true
 
 RSpec.describe "proxying original method params" do # rubocop:disable RSpec/DescribeClass
-  describe ".instance_method" do
+  shared_examples ".instance_method proxies parameters" do
     subject { unbound_method&.parameters }
 
-    let(:unbound_method) { class_with_memo.instance_method(method_name) }
-
-    let(:class_with_memo) do
-      Class.new do
-        prepend MemoWise
-
-        def initialize(foo, bar:); end # rubocop:disable Style/RedundantInitialize
-
-        DefineMethodsForTestingMemoWise.define_methods_for_testing_memo_wise(
-          target: self,
-          via: :instance
-        )
-
-        def unmemoized_with_positional_and_keyword_args(a, b:) # rubocop:disable Naming/MethodParameterName
-          [a, b]
-        end
-      end
-    end
+    let(:unbound_method) { target.instance_method(method_name) }
 
     context "when #initialize" do
       let(:method_name) { :initialize }
@@ -86,4 +69,69 @@ def unmemoized_with_positional_and_keyword_args(a, b:) # rubocop:disable Naming/
       end
     end
   end
+
+  context "when class prepends MemoWise" do
+    let(:target) do
+      Class.new do
+        prepend MemoWise
+
+        def initialize(foo, bar:); end
+
+        DefineMethodsForTestingMemoWise.define_methods_for_testing_memo_wise(
+          target: self,
+          via: :instance
+        )
+
+        def unmemoized_with_positional_and_keyword_args(a, b:) # rubocop:disable Naming/MethodParameterName
+          [a, b]
+        end
+      end
+    end
+
+    it_behaves_like ".instance_method proxies parameters"
+  end
+
+  context "module which prepends MemoWise" do
+    let(:module1) do
+      Module.new do
+        prepend MemoWise
+
+        def initialize(foo, bar:); end
+
+        DefineMethodsForTestingMemoWise.define_methods_for_testing_memo_wise(
+          target: self,
+          via: :instance
+        )
+
+        def unmemoized_with_positional_and_keyword_args(a, b:) # rubocop:disable Naming/MethodParameterName
+          [a, b]
+        end
+      end
+    end
+
+    before(:each) { stub_const("Module1", module1) }
+
+    context "when class includes module" do
+      let(:target) do
+        Class.new do
+          include Module1
+          def initialize(foo, bar:); end
+        end
+      end
+
+      it_behaves_like ".instance_method proxies parameters"
+    end
+
+    context "when class prepends MemoWise and includes module" do
+      let(:target) do
+        Class.new do
+          prepend MemoWise
+          include Module1
+          def initialize(foo, bar:); end
+        end
+      end
+
+      it_behaves_like ".instance_method proxies parameters"
+    end
+  end
 end