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.

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,32 @@ 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
+  end
+
 end