Repository: adobe/elixir-styler
Pull request #9 (2023-04-07 20:44:10)

Change `Style.run` type spec to explicitly rely on `Zipper.traverse_while/3`
zipper traversal is now powered by `traverse_while/3`, and all styles are forced to acknowledge the new shape - no more hiding behind `wrap_run`

```elixir
  @callback run(Zipper.zipper(), context()) :: {Zipper.command(), Zipper.zipper(), context()}
```

future features are likely to work with `comments`; some may want the filename (module directives, for example) and who knows what else we'll add in. passing a map accumulator of context or state gives us a decently future-proof api


Patch:
diff --git a/lib/mix/tasks/style.ex b/lib/mix/tasks/style.ex
index f13a012..ea80869 100644
--- a/lib/mix/tasks/style.ex
+++ b/lib/mix/tasks/style.ex
@@ -30,7 +30,6 @@ defmodule Mix.Tasks.Style do
 
   use Mix.Task
 
-  alias Styler.Style
   alias Styler.StyleError
   alias Styler.Zipper
 
@@ -42,8 +41,6 @@ defmodule Mix.Tasks.Style do
 
   @impl Mix.Task
   def run(args) do
-    for style <- @styles, do: Code.ensure_loaded!(style)
-
     # we take `check_formatted` so we can easily replace `mix format`
     {opts, files} = OptionParser.parse!(args, strict: [check_styled: :boolean, check_formatted: :boolean])
     check_styled? = opts[:check_styled] || opts[:check_formatted] || false
@@ -107,28 +104,29 @@ defmodule Mix.Tasks.Style do
 
     {ast, comments} = Styler.string_to_quoted_with_comments(input, to_string(file))
     zipper = Zipper.zip(ast)
+    context = %{comments: comments, file: file}
 
-    output =
-      @styles
-      |> Enum.reduce(zipper, fn style, zipper ->
-        traverser = Style.wrap_run(style)
-
+    {zipper, %{comments: comments}} =
+      Enum.reduce(@styles, {zipper, context}, fn style, {zipper, context} ->
         try do
-          Zipper.traverse_while(zipper, traverser)
+          Zipper.traverse_while(zipper, context, &style.run/2)
         rescue
           exception ->
             reraise StyleError, [exception: exception, style: style, file: file], __STACKTRACE__
         end
       end)
+
+    styled =
+      zipper
       |> Zipper.root()
       |> Styler.quoted_to_string(comments, formatter_opts)
 
-    changed? = input != output
+    changed? = input != styled
 
     cond do
       check_styled? and changed? -> {:not_styled, file}
-      file == :stdin -> IO.write([output, "\n"])
-      changed? -> File.write!(file, [output, "\n"])
+      file == :stdin -> IO.write([styled, "\n"])
+      changed? -> File.write!(file, [styled, "\n"])
       true -> :ok
     end
   rescue
diff --git a/lib/style.ex b/lib/style.ex
index 3fc5580..1329b04 100644
--- a/lib/style.ex
+++ b/lib/style.ex
@@ -18,7 +18,10 @@ defmodule Styler.Style do
 
   alias Styler.Zipper
 
-  @type command :: :cont | :skip | :halt
+  @type context :: %{
+          comment: [map()],
+          file: :stdin | String.t()
+        }
 
   @doc """
   `run` will be used with `Zipper.traverse_while/3`, meaning it will be executed on every node of the AST.
@@ -26,17 +29,5 @@ defmodule Styler.Style do
   You can skip traversing parts of the tree by returning a Zipper that's further along in the traversal, for example
   by calling `Zipper.skip(zipper)` to skip an entire subtree you know is of no interest to your Style.
   """
-  @callback run(Zipper.zipper()) :: Zipper.zipper() | {command(), Zipper.zipper()}
-
-  @doc false
-  # this lets Styles optionally implement as though they're running inside of `Zipper.traverse`
-  # or `Zipper.traverse_while` for finer-grained control
-  def wrap_run(style) do
-    fn zipper ->
-      case style.run(zipper) do
-        {next, {_, _} = _zipper} = command when next in ~w(cont halt skip)a -> command
-        zipper -> {:cont, zipper}
-      end
-    end
-  end
+  @callback run(Zipper.zipper(), context()) :: {Zipper.command(), Zipper.zipper(), context()}
 end
diff --git a/lib/style/defs.ex b/lib/style/defs.ex
index 6a4b9d0..0cf762e 100644
--- a/lib/style/defs.ex
+++ b/lib/style/defs.ex
@@ -47,20 +47,20 @@ defmodule Styler.Style.Defs do
   #
   #  def example(foo, bar \\ nil)
   #
-  def run({{def, meta, [head]}, _} = zipper) when def in [:def, :defp] do
+  def run({{def, meta, [head]}, _} = zipper, ctx) when def in [:def, :defp] do
     # There won't be any defs deeper in here, so lets skip ahead if we can
-    {:skip, Zipper.replace(zipper, {def, meta, [flatten_head(head, meta[:line])]})}
+    {:skip, Zipper.replace(zipper, {def, meta, [flatten_head(head, meta[:line])]}), ctx}
   end
 
   # all the other kinds of defs!
-  def run({{def, def_meta, [head, body]}, _} = zipper) when def in [:def, :defp] do
+  def run({{def, def_meta, [head, body]}, _} = zipper, ctx) when def in [:def, :defp] do
     def_start_line = def_meta[:line]
     # order matters here! the end of the def is where the `do`s line is - but only if there's a do end block.
     # otherwise it's just where the end of the (def) expression is.
     def_end_line = (def_meta[:do] || def_meta[:end_of_expression])[:line]
 
     if def_start_line == def_end_line do
-      {:skip, zipper}
+      {:skip, zipper, ctx}
     else
       head = flatten_head(head, def_start_line)
 
@@ -87,11 +87,12 @@ defmodule Styler.Style.Defs do
       body = update_all_meta(body, body_meta_rewriter)
 
       # There won't be any defs deeper in here, so lets skip ahead if we can
-      {:skip, Zipper.replace(zipper, {def, def_meta, [head, body]})}
+      # @TODO this skips checking the body, which can be incorrect if therey's a `quote do def do ...` inside of it
+      {:skip, Zipper.replace(zipper, {def, def_meta, [head, body]}), ctx}
     end
   end
 
-  def run(zipper), do: zipper
+  def run(zipper, ctx), do: {:cont, zipper, ctx}
 
   defp collapse_lines(line_mover) do
     fn meta ->
diff --git a/lib/style/module_directives.ex b/lib/style/module_directives.ex
index 8997156..1fdbd62 100644
--- a/lib/style/module_directives.ex
+++ b/lib/style/module_directives.ex
@@ -72,7 +72,7 @@ defmodule Styler.Style.ModuleDirectives do
   @dont_moduledoc ~w(Test Mixfile MixProject Controller Endpoint Repo Router Socket View HTML JSON)
   @moduledoc_false {:@, [], [{:moduledoc, [], [{:__block__, [], [false]}]}]}
 
-  def run({{:defmodule, def_meta, [name, [{mod_do, {:__block__, children_meta, children}}]]}, zipper_meta}) do
+  def run({{:defmodule, def_meta, [name, [{mod_do, {:__block__, children_meta, children}}]]}, zipper_meta}, ctx) do
     {directives, other} =
       Enum.split_with(children, fn
         {:@, _, [{attr, _, _}]} -> attr in @attr_directives
@@ -112,34 +112,38 @@ defmodule Styler.Style.ModuleDirectives do
         other
       ])
 
-    {{:defmodule, def_meta, [name, [{mod_do, {:__block__, children_meta, children}}]]}, zipper_meta}
+    zipper = {{:defmodule, def_meta, [name, [{mod_do, {:__block__, children_meta, children}}]]}, zipper_meta}
+    {:cont, zipper, ctx}
   end
 
   # a module whose only child is a moduledoc. pass it on through
-  def run({{:defmodule, _, [_, [{_, {:@, _, [{:moduledoc, _, _}]}}]]}, _} = zipper), do: zipper
+  def run({{:defmodule, _, [_, [{_, {:@, _, [{:moduledoc, _, _}]}}]]}, _} = zipper, ctx), do: {:cont, zipper, ctx}
 
-  def run({{:defmodule, def_meta, [name, [{mod_do, mod_children}]]}, zipper_meta} = zipper) do
+  def run({{:defmodule, def_meta, [name, [{mod_do, mod_children}]]}, zipper_meta} = zipper, ctx) do
     # a module with a single child. lets add moduledoc false
     # ... unless it's a `defmodule Foo, do: ...`, that is
     if needs_moduledoc?(name, mod_do) do
       mod_children = {:__block__, [], [@moduledoc_false, mod_children]}
-      {{:defmodule, def_meta, [name, [{mod_do, mod_children}]]}, zipper_meta}
+      {:cont, {{:defmodule, def_meta, [name, [{mod_do, mod_children}]]}, zipper_meta}, ctx}
     else
-      zipper
+      {:cont, zipper, ctx}
     end
   end
 
-  def run({{:use, _, _} = directive, meta}) do
+  def run({{:use, _, _} = directive, meta}, ctx) do
     [last | rest] = directive |> expand_directive() |> Enum.reverse()
     meta = %{meta | l: rest ++ meta.l}
 
-    case meta.r do
-      [{:use, _, _} | _] -> {last, meta}
-      _ -> {set_newlines(last, 2), meta}
-    end
+    zipper =
+      case meta.r do
+        [{:use, _, _} | _] -> {last, meta}
+        _ -> {set_newlines(last, 2), meta}
+      end
+
+    {:skip, zipper, ctx}
   end
 
-  def run({{d, _, _} = directive, %{l: left, r: right} = meta}) when d in @directives do
+  def run({{d, _, _} = directive, %{l: left, r: right} = meta}, ctx) when d in @directives do
     {right, directives} = consume_directive_group(d, [directive | right], [])
 
     [last | rest] =
@@ -151,10 +155,12 @@ defmodule Styler.Style.ModuleDirectives do
       |> List.keysort(1, :desc)
       |> Enum.map(&(&1 |> elem(0) |> set_newlines(1)))
 
-    {set_newlines(last, 2), %{meta | r: right, l: rest ++ left}}
+    zipper = {set_newlines(last, 2), %{meta | r: right, l: rest ++ left}}
+
+    {:skip, zipper, ctx}
   end
 
-  def run(zipper), do: zipper
+  def run(zipper, ctx), do: {:cont, zipper, ctx}
 
   def needs_moduledoc?({_, _, aliases}) do
     name = aliases |> List.last() |> to_string()
diff --git a/lib/style/pipes.ex b/lib/style/pipes.ex
index f7db156..d80f311 100644
--- a/lib/style/pipes.ex
+++ b/lib/style/pipes.ex
@@ -26,12 +26,12 @@ defmodule Styler.Style.Pipes do
   @blocks ~w(case if with cond for unless)a
 
   # we're in a multi-pipe, so only need to fix pipe_start
-  def run({{:|>, _, [{:|>, _, _} | _]}, _} = zipper), do: zipper |> check_start() |> Zipper.next()
+  def run({{:|>, _, [{:|>, _, _} | _]}, _} = zipper, ctx), do: {:cont, zipper |> check_start() |> Zipper.next(), ctx}
   # this is a single pipe, since valid pipelines are consumed by the previous head
-  def run({{:|>, meta, [lhs, {fun, _, args}]}, _} = zipper) do
+  def run({{:|>, meta, [lhs, {fun, _, args}]}, _} = zipper, ctx) do
     if valid_pipe_start?(lhs) do
       # `a |> f(b, c)` => `f(a, b, c)`
-      Zipper.replace(zipper, {fun, meta, [lhs | args]})
+      {:cont, Zipper.replace(zipper, {fun, meta, [lhs | args]}), ctx}
     else
       zipper = fix_start(zipper)
       {maybe_block, _, _} = lhs
@@ -39,16 +39,16 @@ defmodule Styler.Style.Pipes do
       if maybe_block in @blocks do
         # extracting a block means this is now `if_result |> single_pipe(a, b)`
         # recursing will give us `single_pipe(if_result, a, b)`
-        run(zipper)
+        run(zipper, ctx)
       else
         # fixing the start when it was a function call added another pipe to the chain, and so it's no longer
         # a single pipe
-        zipper
+        {:cont, zipper, ctx}
       end
     end
   end
 
-  def run(zipper), do: zipper
+  def run(zipper, ctx), do: {:cont, zipper, ctx}
 
   # walking down a pipeline.
   # for reference, `a |> b() |> c()` is encoded `{:|>, [{:|>, _, [a, b]}, c]}`
diff --git a/lib/style/simple.ex b/lib/style/simple.ex
index 6956253..cbf03b7 100644
--- a/lib/style/simple.ex
+++ b/lib/style/simple.ex
@@ -23,7 +23,7 @@ defmodule Styler.Style.Simple do
   alias Styler.Zipper
 
   # `?-` isn't part of the number node - it's its parent - so all numbers are positive at this point
-  def run({{:__block__, meta, [number]}, _} = zipper) when is_number(number) and number >= 10_000 do
+  def run({{:__block__, meta, [number]}, _} = zipper, ctx) when is_number(number) and number >= 10_000 do
     # Checking here rather than in the anon function due to compiler bug https://github.com/elixir-lang/elixir/issues/10485
     integer? = is_integer(number)
 
@@ -47,10 +47,10 @@ defmodule Styler.Style.Simple do
           "#{delimit(int_token)}.#{decimals}"
       end)
 
-    Zipper.replace(zipper, {:__block__, meta, [number]})
+    {:skip, Zipper.replace(zipper, {:__block__, meta, [number]}), ctx}
   end
 
-  def run(zipper), do: zipper
+  def run(zipper, ctx), do: {:cont, zipper, ctx}
 
   defp delimit(token), do: token |> String.to_charlist() |> remove_underscores([]) |> add_underscores([])
 
diff --git a/lib/zipper.ex b/lib/zipper.ex
index db22daa..5e77d5d 100644
--- a/lib/zipper.ex
+++ b/lib/zipper.ex
@@ -34,6 +34,7 @@ defmodule Styler.Zipper do
           }
 
   @type zipper :: {tree, path | nil}
+  @type command :: :cont | :skip | :halt
 
   @doc """
   Returns true if the node is a branch.
@@ -313,7 +314,6 @@ defmodule Styler.Zipper do
   The function must return a zipper.
   """
   @spec traverse_while(zipper, (zipper -> {command, zipper})) :: zipper
-        when command: :cont | :halt | :skip
   def traverse_while({_tree, nil} = zipper, fun) do
     do_traverse_while(zipper, fun)
   end
@@ -341,7 +341,6 @@ defmodule Styler.Zipper do
   If the zipper is not at the top, just the subtree will be traversed.
   """
   @spec traverse_while(zipper, term, (zipper, term -> {command, zipper, term})) :: {zipper, term}
-        when command: :cont | :halt | :skip
   def traverse_while({_tree, nil} = zipper, acc, fun) do
     do_traverse_while(zipper, acc, fun)
   end


Test patch:
diff --git a/test/support/style_case.ex b/test/support/style_case.ex
index 635f1b3..22af6b4 100644
--- a/test/support/style_case.ex
+++ b/test/support/style_case.ex
@@ -14,7 +14,6 @@ defmodule Styler.StyleCase do
   """
   use ExUnit.CaseTemplate
 
-  alias Styler.Style
   alias Styler.Zipper
 
   using options do
@@ -58,14 +57,13 @@ defmodule Styler.StyleCase do
   def style(code, style) do
     {ast, comments} = Styler.string_to_quoted_with_comments(code)
 
-    styled_ast =
+    {zipper, %{comments: comments}} =
       ast
       |> Zipper.zip()
-      |> Zipper.traverse_while(Style.wrap_run(style))
-      |> Zipper.root()
+      |> Zipper.traverse_while(%{comments: comments, file: "test"}, &style.run/2)
 
+    styled_ast = Zipper.root(zipper)
     styled_code = Styler.quoted_to_string(styled_ast, comments)
-
     {styled_ast, styled_code}
   end
 end