Memoize errors

Now errors do not prevent memoization. If the method is called with the same arguments, the same Exception object is raised without an additional call to the method.
matthewrudy · Apr 9, 2015 · 94b88d1 · 94b88d1
1 parent 7c63fe5
commit 94b88d1
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 17 deletions.
diff --git a/lib/memoist.rb b/lib/memoist.rb
@@ -83,6 +83,22 @@ def flush_cache(*method_names)
     end
   end
 
+  class Result
+    attr_reader :error, :value
+
+    def initialize
+      begin
+        @value = yield
+      rescue => e
+        @error = e
+      end
+    end
+
+    def fetch
+      error ? raise(error) : value
+    end
+  end
+
   def memoize(*method_names)
     if method_names.last.is_a?(Hash)
       identifier = method_names.pop[:identifier]
@@ -109,16 +125,18 @@ def memoize(*method_names)
           #   set_cache = skip_cache && !frozen?
           #
           #   if skip_cache
-          #     value = _unmemoized_mime_type
+          #     result = ::Memoist::Result.new do
+          #       _unmemoized_mime_type
+          #     end
           #   else
-          #     value = @_memoized_mime_type
+          #     result = @_memoized_mime_type
           #   end
           #
           #   if set_cache
-          #     @_memoized_mime_type = value
+          #     @_memoized_mime_type = result
           #   end
           #
-          #   value
+          #   result.fetch
           # end
 
           module_eval <<-EOS, __FILE__, __LINE__ + 1
@@ -127,16 +145,18 @@ def #{method_name}(reload = false)
               set_cache = skip_cache && !frozen?
 
               if skip_cache
-                value = #{unmemoized_method}
+                result = ::Memoist::Result.new do
+                  #{unmemoized_method}
+                end
               else
-                value = #{memoized_ivar}
+                result = #{memoized_ivar}
               end
 
               if set_cache
-                #{memoized_ivar} = value
+                #{memoized_ivar} = result
               end
 
-              value
+              result.fetch
             end
           EOS
         else
@@ -150,17 +170,19 @@ def #{method_name}(reload = false)
           #   set_cache = skip_cache && !frozen
           #
           #   if skip_cache
-          #     value = _unmemoized_mime_type(*args)
+          #     result = ::Memoist::Result.new do
+          #       _unmemoized_mime_type(*args)
+          #     end
           #   else
-          #     value = @_memoized_mime_type[args]
+          #     result = @_memoized_mime_type[args]
           #   end
           #
           #   if set_cache
           #     @_memoized_mime_type ||= {}
-          #     @_memoized_mime_type[args] = value
+          #     @_memoized_mime_type[args] = result
           #   end
           #
-          #   value
+          #   result.fetch
           # end
 
           module_eval <<-EOS, __FILE__, __LINE__ + 1
@@ -171,17 +193,19 @@ def #{method_name}(*args)
               set_cache = skip_cache && !frozen?
 
               if skip_cache
-                value = #{unmemoized_method}(*args)
+                result = ::Memoist::Result.new do
+                  #{unmemoized_method}(*args)
+                end
               else
-                value = #{memoized_ivar}[args]
+                result = #{memoized_ivar}[args]
               end
 
               if set_cache
                 #{memoized_ivar} ||= {}
-                #{memoized_ivar}[args] = value
+                #{memoized_ivar}[args] = result
               end
 
-              value
+              result.fetch
             end
           EOS
         end

diff --git a/test/memoist_test.rb b/test/memoist_test.rb
@@ -142,9 +142,10 @@ class Calculator
     extend Memoist
     include Rates
 
-    attr_reader :fib_calls
+    attr_reader :fib_calls, :divide_by_zero_calls
     def initialize
       @fib_calls = 0
+      @divide_by_zero_calls = 0
     end
 
     def fib(n)
@@ -172,6 +173,12 @@ def counter
       @count += 1
     end
     memoize :counter
+
+    def divide_by_zero(x)
+      @divide_by_zero_calls += 1
+      x / 0
+    end
+    memoize :divide_by_zero
   end
 
   def setup
@@ -380,4 +387,38 @@ def test_private_method_memoization
     assert_equal 1, person.is_developer_calls
   end
 
+  def test_memoization_with_an_error
+    assert_equal 0, @calculator.divide_by_zero_calls
+
+    error123 = assert_raises(ZeroDivisionError) do
+      @calculator.divide_by_zero 123
+    end
+    assert_equal 1, @calculator.divide_by_zero_calls
+
+    e = assert_raises(ZeroDivisionError) do
+      @calculator.divide_by_zero 123
+    end
+    assert_same error123, e
+    assert_equal 1, @calculator.divide_by_zero_calls
+
+    error456 = assert_raises(ZeroDivisionError) do
+      @calculator.divide_by_zero 456
+    end
+    assert_equal 2, @calculator.divide_by_zero_calls
+
+    e = assert_raises(ZeroDivisionError) do
+      @calculator.divide_by_zero 456
+    end
+    assert_same error456, e
+    assert_equal 2, @calculator.divide_by_zero_calls
+
+    refute_same error123, error456
+
+    e = assert_raises(ZeroDivisionError) do
+      @calculator.divide_by_zero 456, :reload
+    end
+    refute_same error456, e
+    assert_equal 3, @calculator.divide_by_zero_calls
+  end
+
 end