From 57478f118818e58e92cf22a90bd9d3e2bcc4257d Mon Sep 17 00:00:00 2001
From: Ky Bishop <kybishop@gmail.com>
Date: Thu, 20 Feb 2025 16:20:08 -0500
Subject: [PATCH] feature: specify piped function exclusions

---
 README.md                 |  4 +++-
 lib/quokka/config.ex      |  5 +++++
 lib/style/pipes.ex        | 27 +++++++++++++++++++++++----
 test/style/pipes_test.exs | 13 +++++++++++++
 4 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 7fe8c9e..998eb11 100644
--- a/README.md
+++ b/README.md
@@ -91,7 +91,8 @@ in `.formatter.exs` to fine tune your setup:
       | :module_directives
       | :pipes
       | :single_node
-    ]
+    ],
+    piped_function_exclusions: [:subquery, :"Repo.update", ...]
   ]
 ]
 ```
@@ -102,6 +103,7 @@ in `.formatter.exs` to fine tune your setup:
 | `:only`                          | Only include the given modules. The special `:line_length` option excludes all changes except line length fixups.                                                                                          | `[]` (all modules included)                                         |
 | `:exclude`                       | Exclude the given modules. This is just a convenience function that filters from the `:only` list.                                                                                                         | `[]` (all modules included)                                         |
 | `:inefficient_function_rewrites` | Rewrite inefficient functions to more efficient form                                                                                                                                                       | `true`                                                              |
+| `:piped_function_exclusions`     | Allows you to specify certain functions that won't be rewritten into a pipe. Particularly good for things like Ecto's `subquery` macro.                                                                    | `[]`                                                                |
 
 ## Credo inspired rewrites
 
diff --git a/lib/quokka/config.ex b/lib/quokka/config.ex
index de3a579..c2ec0fe 100644
--- a/lib/quokka/config.ex
+++ b/lib/quokka/config.ex
@@ -92,6 +92,7 @@ defmodule Quokka.Config do
         pipe_chain_start_excluded_argument_types: credo_opts[:pipe_chain_start_excluded_argument_types] || [],
         pipe_chain_start_excluded_functions: credo_opts[:pipe_chain_start_excluded_functions] || [],
         pipe_chain_start_flag: credo_opts[:pipe_chain_start_flag] || false,
+        piped_function_exclusions: Keyword.get(config, :piped_function_exclusions, []),
         rewrite_multi_alias: credo_opts[:rewrite_multi_alias] || false,
         single_pipe_flag: credo_opts[:single_pipe_flag] || false,
         sort_order: credo_opts[:sort_order] || :alpha,
@@ -206,6 +207,10 @@ defmodule Quokka.Config do
     get(:strict_module_layout_order)
   end
 
+  def piped_function_exclusions() do
+    get(:piped_function_exclusions)
+  end
+
   def zero_arity_parens?() do
     get(:zero_arity_parens)
   end
diff --git a/lib/style/pipes.ex b/lib/style/pipes.ex
index ca09375..21f2059 100644
--- a/lib/style/pipes.ex
+++ b/lib/style/pipes.ex
@@ -95,10 +95,10 @@ defmodule Quokka.Style.Pipes do
                 # Not going to lie, no idea why the `shift + 1` is correct but it makes tests pass ¯\_(ツ)_/¯
                 rhs_max_line = Style.max_line(rhs)
 
-              comments =
-                ctx.comments
-                |> Style.displace_comments(lhs_line..(rhs_line - 1)//1)
-                |> Style.shift_comments(rhs_line..rhs_max_line, shift + 1)
+                comments =
+                  ctx.comments
+                  |> Style.displace_comments(lhs_line..(rhs_line - 1)//1)
+                  |> Style.shift_comments(rhs_line..rhs_max_line, shift + 1)
 
                 {:cont, Zipper.replace(single_pipe_zipper, {fun, meta, [lhs | args]}), %{ctx | comments: comments}}
               else
@@ -168,6 +168,14 @@ defmodule Quokka.Style.Pipes do
       function_name in [:assert, :refute | @special_ops] ->
         {:cont, zipper, ctx}
 
+      # Ignore explicitly excluded functions
+      function_name in Quokka.Config.piped_function_exclusions() ->
+        {:cont, zipper, ctx}
+
+      # Ignore explicitly included functions that are part of another module, e.g. Repo.update
+      alias_function_usage_to_existing_atom(function_name) in Quokka.Config.piped_function_exclusions() ->
+        {:cont, zipper, ctx}
+
       # if a |> b() |> c(), do: ...
       Enum.any?(args, &Style.do_block?/1) ->
         {:cont, zipper, ctx}
@@ -183,6 +191,17 @@ defmodule Quokka.Style.Pipes do
 
   def run(zipper, ctx), do: {:cont, zipper, ctx}
 
+  # Functions should look like this:     # {:., [line: 1], [{:__aliases__, [last: [line: 1], line: 1], [:Repo]}, :update]}
+  defp alias_function_usage_to_existing_atom(
+         {:., _metadata, [{:__aliases__, _more_metadata, modules}, function_name]} = _node
+       ) do
+    String.to_existing_atom("#{Enum.join(modules, ".")}.#{function_name}")
+  rescue
+    _ -> nil
+  end
+
+  defp alias_function_usage_to_existing_atom(_), do: nil
+
   defp fix_pipe_start({pipe, zmeta} = zipper) do
     {{:|>, pipe_meta, [lhs, rhs]}, _} = start_zipper = find_pipe_start({pipe, nil})
 
diff --git a/test/style/pipes_test.exs b/test/style/pipes_test.exs
index a5c5df6..82c12b6 100644
--- a/test/style/pipes_test.exs
+++ b/test/style/pipes_test.exs
@@ -1044,6 +1044,19 @@ defmodule Quokka.Style.PipesTest do
         )
       end
     end
+
+    test "doesn't rewrite pipe start functions that are a part of piped_function_exclusions" do
+      # Sanity check
+      assert_style("foo(bar() |> baz() |> boz())", "bar() |> baz() |> boz() |> foo()")
+
+      Quokka.Config.set_for_test!(:piped_function_exclusions, [:foo, :"Ecto.Repo.update"])
+
+      assert_style("foo(bar() |> baz() |> boz())")
+
+      assert_style("Ecto.Repo.update(bar() |> baz() |> boz())")
+
+      Quokka.Config.set_for_test!(:piped_function_exclusions, [])
+    end
   end
 
   describe "comments and..." do