Add parenthesis when an await is within an await or call

Author: bryanjos

lib/es_tree/tools/generator.ex

@@ -263,6 +263,10 @@ defmodule ESTree.Tools.Generator do
 
   # AwaitExpression
 
+  defp do_generate(%AwaitExpression{argument: %AwaitExpression{} = argument, all: _all}, opts) do
+    ["await ", "(", do_generate(argument, opts), ")"]
+  end
+
   defp do_generate(%AwaitExpression{argument: argument, all: _all}, opts) do
     ["await ", do_generate(argument, opts)]
   end
@@ -347,7 +351,12 @@ defmodule ESTree.Tools.Generator do
     end
 
     arguments = arguments
-    |> Enum.map(&do_generate(&1, opts))
+    |> Enum.map(fn
+        %s{} = x when s in [AwaitExpression, YieldExpression] ->
+          ["("] ++ do_generate(x, opts) ++ [")"]
+        x ->
+          do_generate(x, opts)
+       end)
     |> Enum.intersperse(opts.comma_sep)
 
     [callee, "(", arguments, ")"]

test/tools/generator/await_expression_test.exs

@@ -0,0 +1,24 @@
+defmodule ESTree.Tools.Generator.AwaitExpression.Test do
+  use ShouldI
+
+  alias ESTree.Tools.Builder
+  import ESTree.Test.Support
+
+  should "convert await expression" do
+    ast = Builder.await_expression(
+      Builder.literal(1)
+    )
+
+    assert_gen ast, "await 1"
+  end
+
+  should "convert await expression that contains an await" do
+    ast = Builder.await_expression(
+      Builder.await_expression(
+        Builder.literal(1)
+      )
+    )
+
+    assert_gen ast, "await (await 1)"
+  end
+end

test/tools/generator/call_expression_test.exs

@@ -91,4 +91,32 @@ defmodule ESTree.Tools.Generator.CallExpression.Test do
     assert_gen ast, "(a + b).x()"
     assert_gen ast, "(a+b).x()", beauty: false
   end
+
+  should "convert call expression with one argument that is an await" do
+    ast = Builder.call_expression(
+      Builder.identifier(:x),
+      [
+        Builder.await_expression(
+          Builder.identifier(:a)
+        )
+      ]
+    );
+
+    assert_gen ast, "x((await a))"
+    assert_gen ast, "x((await a))", beauty: false
+  end
+
+  should "convert call expression with one argument that is a yield" do
+    ast = Builder.call_expression(
+      Builder.identifier(:x),
+      [
+        Builder.yield_expression(
+          Builder.identifier(:a)
+        )
+      ]
+    );
+
+    assert_gen ast, "x((yield a))"
+    assert_gen ast, "x((yield a))", beauty: false
+  end
 end