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,7 +9,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 **Gem enhancements:** none
 
-_No breaking changes!_
+- Allow memoization when `include`d modules `prepend MemoWise` and define `initialize` [[#327]](https://github.com/panorama-ed/memo_wise/pull/327)
 
 **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.freeze
     # 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,15 @@ def inherited(subclass)
   end
   private_constant(:CreateMemoWiseStateOnInherited)
 
+  module CreateMemoWiseStateOnIncluded
+    def included(base)
+      base.prepend(Module.new do
+        class_eval(INITIALIZE_LITERAL, __FILE__, __LINE__ + 1)
+      end)
+    end
+  end
+  private_constant(:CreateMemoWiseStateOnIncluded)
+
   # @private
   #
   # Private setup method, called automatically by `prepend MemoWise` in a class.
@@ -155,8 +166,7 @@ def memo_wise(method_name_or_hash)
           end
         when Hash
           unless method_name_or_hash.keys == [:self]
-            raise ArgumentError,
-                  "`:self` is the only key allowed in memo_wise"
+            raise ArgumentError, "`:self` is the only key allowed in memo_wise"
           end
 
           method_name = method_name_or_hash[:self]
@@ -175,6 +185,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