Add :only & :except opts like Kernel.import/2 to trace decorator

lib/sibyl/decorator.ex

@@ -11,12 +11,12 @@
  telemetry events for `:start`, `:end`, and `:exception` sub-events.
 
  - `@decorate_all trace()` which does the same as the above, but automatically for
- all functions in a module.
+ all functions in a module (except functions that starts with `_`).
 
  Automatically traced functions are available for reflection by via `Sibyl.Events.reflect/1`.
  """
 
- use Decorator.Define, trace: 0
+ use Decorator.Define, trace: 0, trace: 1
 
  alias Sibyl.AST
  require Sibyl.Events
@@ -27,6 +27,15 @@
  The name of the captured event will be determined by however `Sibyl` is configured
  to generate event names.
 
+ ## Opts
+
+ - `:only` – will decorate only specified functions;
+ - `:except` – will decorate all functions except specified functions.
+
+ Functions are specified as `{name, arity}` tuples, `arity` can be `:*` to specify any arity.
+
+ ## Notes
+
  Due to how anonymous functions are defined and executed in the BEAM, it ends up
  being quite a bit more performant to build the span manually rather than using
  `:telemetry.span/3`.
@@ -37,8 +46,34 @@
  See [here](https://github.com/beam-telemetry/telemetry/pull/43) for explanations
  w.r.t. anonymous function perf.
  """
- @spec trace(function_body :: AST.ast(), ctx :: map()) :: AST.ast() | no_return()
- def trace(body, ctx) do
+ @spec trace(opts, function_body :: AST.ast(), ctx :: map()) :: AST.ast() | no_return()
+ when fun_arity: {name :: atom(), arity :: non_neg_integer() | :*},
+ opts: {:only | :except, [fun_arity]}
+ def trace(opts \\ [], body, ctx) do
+ if should_decorate?(opts, ctx) do
+ do_decorate(body, ctx)
+ else
+ body
+ end
+ end
+
+ @doc false
+ @spec on_definition(env :: map(), term(), atom(), [term()], AST.ast(), AST.ast()) :: AST.ast()
+ def on_definition(%{module: module} = env, kind, function, args, guards, body) do
+ arity = length(args)
+
+ if has_decorate_all?(module) or has_decorate?(module, function) do
+ for event <- [:start, :stop, :exception] do
+ module
+ |> Sibyl.Events.build_event(function, arity, event)
+ |> Sibyl.Events.define_event(module)
+ end
+ end
+
+ Decorator.Decorate.on_definition(env, kind, function, args, guards, body)
+ end
+
+ defp do_decorate(body, ctx) do
  Application.ensure_all_started(:telemetry)
  event = Sibyl.Events.build_event(ctx.module, ctx.name, ctx.arity)
 
@@ -111,20 +146,14 @@
  end
  end
 
- @doc false
- @spec on_definition(env :: map(), term(), atom(), [term()], AST.ast(), AST.ast()) :: AST.ast()
- def on_definition(%{module: module} = env, kind, function, args, guards, body) do
- arity = length(args)
-
- if has_decorate_all?(module) or has_decorate?(module, function) do
- for event <- [:start, :stop, :exception] do
- module
- |> Sibyl.Events.build_event(function, arity, event)
- |> Sibyl.Events.define_event(module)
- end
+ defp should_decorate?(opts, %{name: name, arity: arity}) do
+ cond do
+ opts[:only][name] in [:*, arity] -> true
+ opts[:only] -> false
+ opts[:except][name] in [:*, arity] -> false
+ name |> to_string() |> String.starts_with?("_") -> false
+ true -> true
  end
-
- Decorator.Decorate.on_definition(env, kind, function, args, guards, body)
  end
 
  defp has_decorate_all?(module), do: Module.get_attribute(module, :decorate_all) != []
@@ -134,7 +163,7 @@
  |> Module.get_attribute(:decorated)
  |> Enum.find(fn node ->
  is_tuple(node) && elem(node, 0) in [:def, :defp] && elem(node, 1) == function &&
- {Sibyl.Decorator, :trace, []} in elem(node, 5)
+ Enum.any?(elem(node, 5), &match?({Sibyl.Decorator, :trace, _}, &1))
  end)
  |> is_tuple()
  end