diff --git a/lib/memo_wise.rb b/lib/memo_wise.rb index dace3d6..5f61e36 100644 --- a/lib/memo_wise.rb +++ b/lib/memo_wise.rb @@ -94,6 +94,20 @@ def inherited(subclass) end private_constant(:CreateMemoWiseStateOnInherited) + module CreateMemoWiseStateOnIncluded + def included(base) + return unless base.is_a?(Class) && !base.singleton_class? + + base.prepend(Module.new do + def initialize(...) + MemoWise::InternalAPI.create_memo_wise_state!(self) + super + end + end) + end + end + private_constant(:CreateMemoWiseStateOnIncluded) + # @private # # Private setup method, called automatically by `prepend MemoWise` in a class. @@ -176,6 +190,12 @@ def memo_wise(method_name_or_hash) if klass.is_a?(Class) && !klass.singleton_class? klass.singleton_class.prepend(CreateMemoWiseStateOnInherited) else + if klass.is_a?(Module) && !klass.singleton_class? + klass.singleton_class.prepend(CreateMemoWiseStateOnIncluded) + elsif klass.is_a?(Module) + klass.prepend(CreateMemoWiseStateOnIncluded) + end + klass.prepend(CreateMemoWiseStateOnInherited) end diff --git a/spec/memo_wise_spec.rb b/spec/memo_wise_spec.rb index 2a3f4b9..0fca4ee 100644 --- a/spec/memo_wise_spec.rb +++ b/spec/memo_wise_spec.rb @@ -351,6 +351,13 @@ def module2_method end end + let(:klass_with_initializer) do + Class.new do + include Module1 + def initialize(*); end + end + end + let(:instance) { klass.new } before(:each) do @@ -364,6 +371,16 @@ 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 + end + + it "can reset klass with initializer" do + instance = klass_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