More flexible result matching (#61)

* Make after & else clauses optional for retry macro

* Run formatter

* Support flexible rescue_only parameters

* Support flexible atoms parameters

* Fix extra runs

* Add a few test cases

* Update documentation

* Fix unused variable warning

* Fix example code in docs

Author: nathanalderson

README.md

@@ -32,9 +32,23 @@ Check out the [API reference](https://hexdocs.pm/retry/api-reference.html) for t
 
 The `retry([with: _,] do: _, after: _, else: _)` macro provides a way to retry a block of code on failure with a variety of delay and give up behaviors. By default, the execution of a block is considered a failure if it returns `:error`, `{:error, _}` or raises a runtime error.
 
-An optional list of atoms can be specified in `:atoms` if you need to retry anything other than `:error` or `{:error, _}`, e.g. `retry([with: _, atoms: [:not_ok]], do: _, after: _, else: _)`.
+Both the values and exceptions that will be retried can be customized. To control which values will be retried, provide the `atoms` option. To control which exceptions are retried, provide the `rescue_only` option. For example:
 
-Similarly, an optional list of exceptions can be specified in `:rescue_only` if you need to retry anything other than `RuntimeError`, e.g. `retry([with: _, rescue_only: [CustomError]], do: _, after: _, else: _)`.
+```
+retry with: ..., atoms: [:not_ok], rescue_only: [CustomError] do
+  ...
+end
+```
+
+Both `atoms` and `rescue_only` can accept a number of different types:
+
+* An atom (for example: `:not_okay`, `SomeStruct`, or `CustomError`). In this case, the `do` block will be retried in any of the following cases:
+  * The atom itself is returned
+  * The atom is returned in the first position of a two-tuple (for example, `{:not_okay, _}`)
+  * A struct of that type is returned/raised
+* The special atom `:all`. In this case, all values/exceptions will be retried.
+* A function (for example: `fn val -> String.starts_with?(val, "ok") end`) or partial function (for example: `fn {:error, %SomeStruct{reason: "busy"}} -> true`). The function will be called with the return value and the `do` block will be retried if the function returns a truthy value. If the function returns a falsy value or if no function clause matches, the `do` block will not be retried.
+* A list of any of the above. The `do` block will be retried if any of the items in the list matches.
 
 The `after` block evaluates only when the `do` block returns a valid value before timeout. On the other hand, the `else` block evaluates only when the `do` block remains erroneous after timeout. Both are optional. By default, the `else` clause will return the last erroneous value or re-raise the last exception. The default `after` clause will simply return the last successful value.
 

lib/retry.ex

@@ -103,13 +103,31 @@ defmodule Retry do
   Retry a block of code delaying between each attempt the duration specified by
   the next item in the `with` delay stream.
 
-  If the block returns any of the atoms specified in `atoms`, a retry will be attempted.
-  Other atoms or atom-result tuples will not be retried. If `atoms` is not specified,
-  it defaults to `[:error]`.
+  Both the values and exceptions that will be retried can be customized. To control which values
+  will be retried, provide the `atoms` option. To control which exceptions are retried, provide
+  the `rescue_only` option. For example:
 
-  Similary, if the block raises any of the exceptions specified in `rescue_only`, a retry
-  will be attempted. Other exceptions will not be retried. If `rescue_only` is
-  not specified, it defaults to `[RuntimeError]`.
+  ```
+  retry with: ..., atoms: [:not_ok], rescue_only: [CustomError] do
+    ...
+  end
+  ```
+
+  Both `atoms` and `rescue_only` can accept a number of different types:
+
+  * An atom (for example: `:not_okay`, `SomeStruct`, or `CustomError`). In this case, the `do`
+    block will be retried in any of the following cases:
+    * The atom itself is returned
+    * The atom is returned in the first position of a two-tuple (for example, `{:not_okay, _}`)
+    * A struct of that type is returned/raised
+  * The special atom `:all`. In this case, all values/exceptions will be retried.
+  * A function (for example: `fn val -> String.starts_with?(val, "ok") end`) or partial function
+    (for example: `fn {:error, %SomeStruct{reason: "busy"}} -> true`). The function will be called
+    with the return value and the `do` block will be retried if the function returns a truthy value.
+    If the function returns a falsy value or if no function clause matches, the `do` block
+    will not be retried.
+  * A list of any of the above. The `do` block will be retried if any of the items in the list
+    matches.
 
   The `after` block evaluates only when the `do` block returns a valid value before timeout.
 
@@ -151,7 +169,6 @@ defmodule Retry do
     opts = parse_opts(opts, @retry_meta)
     [do_clause, after_clause, else_clause] = parse_clauses(clauses, @retry_meta)
     stream_builder = Keyword.fetch!(opts, :with)
-    atoms = Keyword.fetch!(opts, :atoms)
 
     quote generated: true do
       fun = unquote(block_runner(do_clause, opts))
@@ -167,12 +184,7 @@ defmodule Retry do
             unquote(else_clause)
           end
 
-        e = {atom, _} when atom in unquote(atoms) ->
-          case e do
-            unquote(else_clause)
-          end
-
-        e when is_atom(e) and e in unquote(atoms) ->
+        {:retriable, e} ->
           case e do
             unquote(else_clause)
           end
@@ -313,29 +325,52 @@ defmodule Retry do
 
   defp block_runner(block, opts) do
     atoms = Keyword.get(opts, :atoms)
-    exceptions = Keyword.get(opts, :rescue_only)
+    rescue_onlies = Keyword.get(opts, :rescue_only)
 
     quote generated: true do
+      call_partial = fn f, x ->
+        try do
+          !!f.(x)
+        rescue
+          FunctionClauseError -> false
+        end
+      end
+
+      should_retry = fn
+        _x, :all -> true
+        x, a when is_atom(x) and is_atom(a) -> x == a
+        x, a when is_struct(x) and is_atom(a) -> is_struct(x, a)
+        {x, _}, a when is_atom(x) and is_atom(a) -> x == a
+        x, f when is_function(f) -> call_partial.(f, x)
+        _, _ -> false
+      end
+
       fn ->
         try do
-          case unquote(block) do
-            {atom, _} = result ->
-              if atom in unquote(atoms) do
-                {:cont, result}
-              else
-                {:halt, result}
-              end
+          result = unquote(block)
 
-            result ->
-              if is_atom(result) and result in unquote(atoms) do
-                {:cont, result}
-              else
-                {:halt, result}
-              end
+          retry? =
+            if is_list(unquote(atoms)) do
+              Enum.any?(unquote(atoms), &should_retry.(result, &1))
+            else
+              should_retry.(result, unquote(atoms))
+            end
+
+          if retry? do
+            {:cont, {:retriable, result}}
+          else
+            {:halt, result}
           end
         rescue
           e ->
-            if e.__struct__ in unquote(exceptions) do
+            retry? =
+              if is_list(unquote(rescue_onlies)) do
+                Enum.any?(unquote(rescue_onlies), &should_retry.(e, &1))
+              else
+                should_retry.(e, unquote(rescue_onlies))
+              end
+
+            if retry? do
               {:cont, {:exception, e}}
             else
               reraise e, __STACKTRACE__

test/retry_test.exs

@@ -3,10 +3,13 @@ defmodule RetryTest do
   use Retry
 
   import Stream
+  import ExUnit.CaptureLog
+  require Logger
 
   doctest Retry
 
   defmodule(CustomError, do: defexception(message: "custom error!"))
+  defmodule(NotOkay, do: defstruct([]))
 
   describe "retry" do
     test "retries execution for specified attempts when result is error tuple" do
@@ -45,44 +48,59 @@ defmodule RetryTest do
       assert elapsed / 1_000 >= 250
     end
 
-    test "retries execution for specified attempts when result is a specified atom" do
-      retry_atom = :not_ok
-
-      {elapsed, _} =
-        :timer.tc(fn ->
-          result =
-            retry with: linear_backoff(50, 1) |> take(5), atoms: [retry_atom] do
-              retry_atom
-            after
-              _ -> :ok
-            else
-              error -> error
-            end
-
-          assert result == retry_atom
-        end)
+    test "retries execution for specified attempts when allowed result is returned" do
+      testcases = [
+        {:not_ok, :all},
+        {:not_ok, [:foo, :all]},
+        {:not_ok, :not_ok},
+        {:not_ok, [:foo, :not_ok]},
+        {{:not_ok, :foo}, [:foo, :not_ok]},
+        {%NotOkay{}, NotOkay},
+        {%NotOkay{}, [Foo, NotOkay]},
+        {:not_ok, fn _ -> true end},
+        {:not_ok, [fn _ -> false end, fn _ -> true end]},
+        {:not_ok, [fn _ -> nil end, fn _ -> 1 end]},
+        {:not_ok, [fn :partial -> false end, fn _ -> true end]},
+        {:not_ok,
+         fn
+           :partial -> false
+           :not_ok -> true
+         end}
+      ]
 
-      assert elapsed / 1_000 >= 250
+      for {rval, atoms} <- testcases do
+        {elapsed, _} =
+          :timer.tc(fn ->
+            result =
+              retry with: linear_backoff(50, 1) |> take(5), atoms: atoms do
+                rval
+              after
+                _ -> :ok
+              else
+                error -> error
+              end
+
+            assert result == rval
+          end)
+
+        assert elapsed / 1_000 >= 250
+      end
     end
 
-    test "retries execution for specified attempts when result is a tuple with a specified atom" do
-      retry_atom = :not_ok
-
-      {elapsed, _} =
-        :timer.tc(fn ->
-          result =
-            retry with: linear_backoff(50, 1) |> take(5), atoms: [retry_atom] do
-              {retry_atom, "Some error message"}
-            after
-              _ -> :ok
-            else
-              error -> error
-            end
-
-          assert result == {retry_atom, "Some error message"}
-        end)
+    test "does not retry on :error if atoms is specified" do
+      f = fn ->
+        retry with: linear_backoff(50, 1) |> take(5), atoms: :not_ok do
+          Logger.info("running")
+          :error
+        after
+          result -> result
+        else
+          _error -> :not_this
+        end
+      end
 
-      assert elapsed / 1_000 >= 250
+      assert f.() == :error
+      assert Regex.scan(~r/running/, capture_log(f)) |> length == 1
     end
 
     test "retries execution for specified attempts when error is raised" do
@@ -102,23 +120,33 @@ defmodule RetryTest do
       assert elapsed / 1_000 >= 250
     end
 
-    test "retries execution when a whitelisted exception is raised" do
-      custom_error_list = [CustomError]
+    test "retries execution when an allowed exception is raised" do
+      testcases = [
+        CustomError,
+        [OtherThing, CustomError],
+        :all,
+        [:other_thing, :all],
+        fn _ -> true end,
+        [fn _ -> false end, fn _ -> true end],
+        [fn :partial -> false end, fn _ -> true end]
+      ]
 
-      {elapsed, _} =
-        :timer.tc(fn ->
-          assert_raise CustomError, fn ->
-            retry with: linear_backoff(50, 1) |> take(5), rescue_only: custom_error_list do
-              raise CustomError
-            after
-              _ -> :ok
-            else
-              error -> raise error
+      for testcase <- testcases do
+        {elapsed, _} =
+          :timer.tc(fn ->
+            assert_raise CustomError, fn ->
+              retry with: linear_backoff(50, 1) |> take(5), rescue_only: testcase do
+                raise CustomError
+              after
+                _ -> :ok
+              else
+                error -> raise error
+              end
             end
-          end
-        end)
+          end)
 
-      assert elapsed / 1_000 >= 250
+        assert elapsed / 1_000 >= 250
+      end
     end
 
     test "does not retry execution when an unknown exception is raised" do
@@ -138,17 +166,37 @@ defmodule RetryTest do
       assert elapsed / 1_000 < 250
     end
 
+    test "does not retry on RuntimeError if some other rescue_only is specified" do
+      f = fn ->
+        assert_raise RuntimeError, fn ->
+          retry with: linear_backoff(50, 1) |> take(5), rescue_only: CustomError do
+            Logger.info("running")
+            raise RuntimeError
+          after
+            _ -> :ok
+          else
+            error -> raise error
+          end
+        end
+      end
+
+      assert Regex.scan(~r/running/, capture_log(f)) |> length == 1
+    end
+
     test "does not have to retry execution when there is no error" do
-      result =
+      f = fn ->
         retry with: linear_backoff(50, 1) |> take(5) do
+          Logger.info("running")
           {:ok, "Everything's so awesome!"}
         after
           result -> result
         else
           _ -> :error
         end
+      end
 
-      assert result == {:ok, "Everything's so awesome!"}
+      assert f.() == {:ok, "Everything's so awesome!"}
+      assert Regex.scan(~r/running/, capture_log(f)) |> length == 1
     end
 
     test "uses the default 'after' action" do