diff --git a/lib/decorators/decorate.ex b/lib/decorators/decorate.ex index 5822e41..f63999d 100644 --- a/lib/decorators/decorate.ex +++ b/lib/decorators/decorate.ex @@ -33,11 +33,10 @@ defmodule Decorator.Decorate do fn({_kind, _fun, _args, _guard, _body, decorators}) -> decorators end) - |> Enum.filter_map( - fn({_k, decorators_list}) -> + |> Enum.filter(fn({_k, decorators_list}) -> List.flatten(decorators_list) != [] - end, - fn({k, _decorators_list}) -> k end) + end) + |> Enum.map(fn({k, _decorators_list}) -> k end) all |> Enum.filter( fn({_kind, fun, args, _guard, _body, _decorators}) -> @@ -72,20 +71,17 @@ defmodule Decorator.Decorate do |> Enum.reduce(body, fn(decorator, body) -> apply_decorator(context, decorator, body) end) + |> ensure_do() def_clause = case guard do [] -> quote do - Kernel.unquote(kind)(unquote(fun)(unquote_splicing(args))) do - unquote(body) - end + Kernel.unquote(kind)(unquote(fun)(unquote_splicing(args)), unquote(body)) end _ -> quote do - Kernel.unquote(kind)(unquote(fun)(unquote_splicing(args)) when unquote_splicing(guard)) do - unquote(body) - end + Kernel.unquote(kind)(unquote(fun)(unquote_splicing(args)) when unquote_splicing(guard), unquote(body)) end end @@ -96,6 +92,15 @@ defmodule Decorator.Decorate do end end + defp ensure_do([{:do, _} | _] = body), do: body + defp ensure_do(body), do: [do: body] + + defp apply_decorator(context, mfa, [do: body]) do + [do: apply_decorator(context, mfa, body)] + end + defp apply_decorator(context, mfa, [do: body, rescue: rescue_block]) do + [do: apply_decorator(context, mfa, body), rescue: apply_decorator_to_rescue(context, mfa, rescue_block)] + end defp apply_decorator(context, {module, fun, args}, body) do if Enum.member?(module.__info__(:exports), {fun, Enum.count(args) + 2}) do Kernel.apply(module, fun, (args || []) ++ [body, context]) @@ -107,6 +112,13 @@ defmodule Decorator.Decorate do raise ArgumentError, "Invalid decorator: #{inspect decorator}" end + defp apply_decorator_to_rescue(context, mfa, rescue_block) do + rescue_block + |> Enum.map(fn({:->, meta, [match, body]}) -> + {:->, meta, [match, apply_decorator(context, mfa, body)]} + end) + end + def generate_args(0, _caller), do: [] def generate_args(n, caller), do: for(i <- 1..n, do: Macro.var(:"var#{i}", caller)) diff --git a/test/exception_test.exs b/test/exception_test.exs new file mode 100644 index 0000000..f6cb65d --- /dev/null +++ b/test/exception_test.exs @@ -0,0 +1,35 @@ +# two arguments +defmodule DecoratorTest.Fixture.ExceptionDecorator do + use Decorator.Define, [test: 0] + + def test(body, _context) do + {:ok, body} + end + +end + +defmodule DecoratorTest.Fixture.ExceptionTestModule do + use DecoratorTest.Fixture.ExceptionDecorator + + @decorate test() + def result(a) do + if a == :throw do + raise RuntimeError, "text" + end + a + rescue + _ in RuntimeError -> + :error + end +end + +defmodule DecoratorTest.Exception do + use ExUnit.Case + + alias DecoratorTest.Fixture.ExceptionTestModule + + test "Functions which have a 'rescue' clause" do + assert {:ok, 3} = ExceptionTestModule.result(3) + assert {:ok, :error} = ExceptionTestModule.result(:throw) + end +end