pipes: treat MapSet.new and Keyword.new the same way we do Map.new

Co-authored-by: Michael Shapiro <[email protected]>
adobe · Jan 25, 2024 · c3c0bc7 · c3c0bc7
1 parent 7db0c25
commit c3c0bc7
Show file tree

Hide file tree

Showing 5 changed files with 131 additions and 113 deletions.
diff --git a/lib/style.ex b/lib/style.ex
@@ -153,9 +153,4 @@ defmodule Styler.Style do
     end)
     |> Enum.sort_by(& &1.line)
   end
-
-  @doc "Returns true if the ast represents an empty map"
-  def empty_map?({:%{}, _, []}), do: true
-  def empty_map?({{:., _, [{_, _, [:Map]}, :new]}, _, []}), do: true
-  def empty_map?(_), do: false
 end
diff --git a/lib/style/pipes.ex b/lib/style/pipes.ex
@@ -29,6 +29,8 @@ defmodule Styler.Style.Pipes do
   alias Styler.Style
   alias Styler.Zipper
 
+  @collectable ~w(Map Keyword MapSet)a
+
   def run({{:|>, _, _}, _} = zipper, ctx) do
     case fix_pipe_start(zipper) do
       {{:|>, _, _}, _} = zipper ->
@@ -189,10 +191,11 @@ defmodule Styler.Style.Pipes do
   defp fix_pipe(
          pipe_chain(
            lhs,
-           {{:., dm, [{_, _, [:Enum]} = enum, :map]}, em, [mapper]},
-           {{:., _, [{_, _, [:Enum]}, :join]}, _, [joiner]}
+           {{:., dm, [{_, _, [enum_or_stream]}, :map]}, em, [mapper]},
+           {{:., _, [{_, _, [:Enum]} = enum, :join]}, _, [joiner]}
          )
-       ) do
+       )
+       when enum_or_stream in [:Enum, :Stream] do
     rhs = Style.set_line({{:., dm, [enum, :map_join]}, em, [joiner, mapper]}, dm[:line])
     {:|>, [line: dm[:line]], [lhs, rhs]}
   end
@@ -203,28 +206,48 @@ defmodule Styler.Style.Pipes do
   defp fix_pipe(
          pipe_chain(
            lhs,
-           {{:., dm, [{_, am, [:Enum]}, :map]}, em, [mapper]},
+           {{:., dm, [{_, _, [mod]}, :map]}, _, [mapper]},
            {{:., _, [{_, _, [:Enum]}, :into]} = into, _, [collectable]}
          )
-       ) do
+       )
+       when mod in [:Enum, :Stream] do
     rhs =
-      if Style.empty_map?(collectable),
-        do: {{:., dm, [{:__aliases__, am, [:Map]}, :new]}, em, [mapper]},
-        else: {into, em, [collectable, mapper]}
+      case collectable do
+        {{:., _, [{_, _, [mod]}, :new]}, _, []} when mod in @collectable ->
+          {{:., dm, [{:__aliases__, dm, [mod]}, :new]}, dm, [mapper]}
+
+        {:%{}, _, []} ->
+          {{:., dm, [{:__aliases__, dm, [:Map]}, :new]}, dm, [mapper]}
+
+        _ ->
+          {into, dm, [collectable, mapper]}
+      end
 
     Style.set_line({:|>, [], [lhs, rhs]}, dm[:line])
   end
 
-  defp fix_pipe({:|>, meta, [lhs, {{:., dm, [{_, am, [:Enum]}, :into]}, em, [collectable | rest]}]} = node) do
-    if Style.empty_map?(collectable),
-      do: {:|>, meta, [lhs, {{:., dm, [{:__aliases__, am, [:Map]}, :new]}, em, rest}]},
-      else: node
+  # lhs |> Enum.into(%{}, ...) => lhs |> Map.new(...)
+  defp fix_pipe({:|>, meta, [lhs, {{:., dm, [{_, _, [:Enum]}, :into]}, _, [collectable | rest]}]} = node) do
+    replacement =
+      case collectable do
+        {{:., _, [{_, _, [mod]}, :new]}, _, []} when mod in @collectable ->
+          {:., dm, [{:__aliases__, dm, [mod]}, :new]}
+
+        {:%{}, _, []} ->
+          {:., dm, [{:__aliases__, dm, [:Map]}, :new]}
+
+        _ ->
+          nil
+      end
+
+    if replacement, do: {:|>, meta, [lhs, {replacement, dm, rest}]}, else: node
   end
 
   # `lhs |> Enum.map(mapper) |> Map.new()` => `lhs |> Map.new(mapper)
   defp fix_pipe(
-         pipe_chain(lhs, {{:., _, [{_, _, [:Enum]}, :map]}, _, [mapper]}, {{:., _, [{_, _, [:Map]}, :new]} = new, nm, []})
-       ) do
+         pipe_chain(lhs, {{:., _, [{_, _, [:Enum]}, :map]}, _, [mapper]}, {{:., _, [{_, _, [mod]}, :new]} = new, nm, []})
+       )
+       when mod in @collectable do
     Style.set_line({:|>, [], [lhs, {new, nm, [mapper]}]}, nm[:line])
   end
 

diff --git a/lib/style/single_node.ex b/lib/style/single_node.ex
@@ -27,7 +27,6 @@ defmodule Styler.Style.SingleNode do
 
   @behaviour Styler.Style
 
-  alias Styler.Style
   alias Styler.Zipper
 
   def run({node, meta}, ctx), do: {:cont, {style(node), meta}, ctx}
@@ -73,13 +72,25 @@ defmodule Styler.Style.SingleNode do
     {:__block__, meta, [number]}
   end
 
-  ## DEPRECATED & INEFFICIENT FUNCTION REWRITES
+  ## INEFFICIENT FUNCTION REWRITES
   # Keep in mind when rewriting a `/n::pos_integer` arity function here that it should also be added
   # to the pipes rewriting rules, where it will appear as `/n-1`
 
   # Enum.into(enum, empty_map[, ...]) => Map.new(enum[, ...])
-  defp style({{:., dm, [{:__aliases__, am, [:Enum]}, :into]}, funm, [enum, collectable | rest]} = node) do
-    if Style.empty_map?(collectable), do: {{:., dm, [{:__aliases__, am, [:Map]}, :new]}, funm, [enum | rest]}, else: node
+  defp style({{:., dm, [{:__aliases__, _, [:Enum]}, :into]}, funm, [enum, collectable | rest]} = node) do
+    new_collectable =
+      case collectable do
+        {{:., _, [{_, _, [mod]}, :new]}, _, []} when mod in ~w(Map Keyword MapSet)a ->
+          {:., dm, [{:__aliases__, dm, [mod]}, :new]}
+
+        {:%{}, _, []} ->
+          {:., dm, [{:__aliases__, dm, [:Map]}, :new]}
+
+        _ ->
+          nil
+      end
+
+    if new_collectable, do: {new_collectable, funm, [enum | rest]}, else: node
   end
 
   for mod <- [:Map, :Keyword] do

diff --git a/test/style/pipes_test.exs b/test/style/pipes_test.exs
@@ -596,87 +596,62 @@ defmodule Styler.Style.PipesTest do
       )
     end
 
-    test "map/into" do
-      assert_style(
-        """
-        a
-        |> Enum.map(b)
-        |> Enum.into(%{})
-        """,
-        "Map.new(a, b)"
-      )
-
-      assert_style(
-        """
-        a
-        |> Enum.map(b)
-        |> Enum.into(Map.new())
-        """,
-        "Map.new(a, b)"
-      )
-
-      assert_style(
-        """
-        a
-        |> Enum.map(b)
-        |> Enum.into(my_map)
-        """,
-        "Enum.into(a, my_map, b)"
-      )
+    test "enum/stream.map |> enum.into" do
+      for enum <- ~w(Enum Stream) do
+        assert_style("a|> #{enum}.map(b)|> Enum.into(%{})", "Map.new(a, b)")
+        assert_style("a |> #{enum}.map(b) |> Enum.into(unk)", "Enum.into(a, unk, b)")
+
+        assert_style(
+          "a |> #{enum}.map(b) |> Enum.into(%{some: :existing_map})",
+          "Enum.into(a, %{some: :existing_map}, b)"
+        )
 
-      assert_style(
-        """
-        a
-        |> Enum.map(b)
-        |> Enum.into(%{some: :existing_map})
-        """,
-        "Enum.into(a, %{some: :existing_map}, b)"
-      )
+        assert_style(
+          """
+          a_multiline_mapper
+          |> #{enum}.map(fn %{gets: shrunk, down: to_a_more_reasonable} ->
+            IO.puts "woo!"
+            {shrunk, to_a_more_reasonable}
+          end)
+          |> Enum.into(size)
+          """,
+          """
+          Enum.into(a_multiline_mapper, size, fn %{gets: shrunk, down: to_a_more_reasonable} ->
+            IO.puts("woo!")
+            {shrunk, to_a_more_reasonable}
+          end)
+          """
+        )
 
-      assert_style(
-        """
-        a_multiline_mapper
-        |> Enum.map(fn %{gets: shrunk, down: to_a_more_reasonable} ->
-          IO.puts "woo!"
-          {shrunk, to_a_more_reasonable}
-        end)
-        |> Enum.into(size)
-        """,
-        """
-        Enum.into(a_multiline_mapper, size, fn %{gets: shrunk, down: to_a_more_reasonable} ->
-          IO.puts("woo!")
-          {shrunk, to_a_more_reasonable}
-        end)
-        """
-      )
+        for collectable <- ~W(Map Keyword MapSet), new = "#{collectable}.new" do
+          assert_style("a |> #{enum}.map(b) |> Enum.into(#{new}())", "#{new}(a, b)")
 
-      # Regression: something about the meta wants newlines when it's in a def
-      assert_style(
-        """
-        def foo() do
-          filename_map = foo |> Enum.map(&{&1.filename, true}) |> Enum.into(%{})
-        end
-        """,
-        """
-        def foo do
-          filename_map = Map.new(foo, &{&1.filename, true})
+          # Regression: something about the meta wants newlines when it's in a def
+          assert_style(
+            """
+            def foo() do
+              filename_map = foo |> Enum.map(&{&1.filename, true}) |> Enum.into(%{})
+            end
+            """,
+            """
+            def foo do
+              filename_map = Map.new(foo, &{&1.filename, true})
+            end
+            """
+          )
         end
-        """
-      )
+      end
     end
 
-    test "Enum.map(x) |> Map.new()" do
-      assert_style("a |> Enum.map(b) |> Map.new()", "Map.new(a, b)")
+    test "Enum.map(x) |> $collectable.new()" do
+      for collectable <- ~W(Map Keyword MapSet), new = "#{collectable}.new" do
+        assert_style("a |> Enum.map(b) |> #{new}()", "#{new}(a, b)")
+      end
     end
 
-    test "into a new map" do
-      assert_style("a |> Enum.into(foo) |> b()")
+    test "into an empty map" do
       assert_style("a |> Enum.into(%{}) |> b()", "a |> Map.new() |> b()")
-      assert_style("a |> Enum.into(Map.new) |> b()", "a |> Map.new() |> b()")
-
-      assert_style("a |> Enum.into(foo, mapper) |> b()")
       assert_style("a |> Enum.into(%{}, mapper) |> b()", "a |> Map.new(mapper) |> b()")
-      assert_style("a |> Enum.into(Map.new, mapper) |> b()", "a |> Map.new(mapper) |> b()")
 
       assert_style(
         """
@@ -690,19 +665,29 @@ defmodule Styler.Style.PipesTest do
         |> Map.new(c)
         """
       )
+    end
 
-      assert_style(
-        """
-        a
-        |> Enum.map(b)
-        |> Enum.into(Map.new, c)
-        """,
-        """
-        a
-        |> Enum.map(b)
-        |> Map.new(c)
-        """
-      )
+    test "into a new collectable" do
+      assert_style("a |> Enum.into(foo) |> b()")
+      assert_style("a |> Enum.into(foo, mapper) |> b()")
+
+      for collectable <- ~W(Map Keyword MapSet), new = "#{collectable}.new" do
+        assert_style("a |> Enum.into(#{new}) |> b()", "a |> #{new}() |> b()")
+        assert_style("a |> Enum.into(#{new}, mapper) |> b()", "a |> #{new}(mapper) |> b()")
+
+        assert_style(
+          """
+          a
+          |> Enum.map(b)
+          |> Enum.into(#{new}, c)
+          """,
+          """
+          a
+          |> Enum.map(b)
+          |> #{new}(c)
+          """
+        )
+      end
     end
   end
 end
diff --git a/test/style/single_node_test.exs b/test/style/single_node_test.exs
@@ -280,16 +280,20 @@ defmodule Styler.Style.SingleNodeTest do
     end
   end
 
-  describe "Enum.into and Map.new" do
-    test "into a new map" do
+  describe "Enum.into and $collectable.new" do
+    test "into an empty map" do
+      assert_style("Enum.into(a, %{})", "Map.new(a)")
+      assert_style("Enum.into(a, %{}, mapper)", "Map.new(a, mapper)")
+    end
+
+    test "into a collectable" do
       assert_style("Enum.into(a, foo)")
       assert_style("Enum.into(a, foo, mapper)")
 
-      assert_style("Enum.into(a, %{})", "Map.new(a)")
-      assert_style("Enum.into(a, Map.new)", "Map.new(a)")
-
-      assert_style("Enum.into(a, %{}, mapper)", "Map.new(a, mapper)")
-      assert_style("Enum.into(a, Map.new, mapper)", "Map.new(a, mapper)")
+      for collectable <- ~W(Map Keyword MapSet), new = "#{collectable}.new" do
+        assert_style("Enum.into(a, #{new})", "#{new}(a)")
+        assert_style("Enum.into(a, #{new}, mapper)", "#{new}(a, mapper)")
+      end
     end
   end