Memoize errors

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

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