From 5d0adab7fbf40a862ecad9eacff86f47c3d9c6ec Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicol=C3=B2=20Ribaudo?= <nicolo.ribaudo@gmail.com>
Date: Wed, 25 Nov 2020 16:51:58 +0100
Subject: [PATCH] Set correct `path.context` un `push/unshiftContainer`
 (#12394)

---
 packages/babel-traverse/src/context.js        |  2 +
 .../babel-traverse/src/path/modification.js   |  4 +-
 packages/babel-traverse/test/modification.js  | 40 +++++++++++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/packages/babel-traverse/src/context.js b/packages/babel-traverse/src/context.js
index 12f76d0c9ad4..be41ecc90c37 100644
--- a/packages/babel-traverse/src/context.js
+++ b/packages/babel-traverse/src/context.js
@@ -42,6 +42,8 @@ export default class TraversalContext {
   }
 
   create(node, obj, key, listKey): NodePath {
+    // We don't need to `.setContext()` here, since `.visitQueue()` already
+    // calls `.pushContext`.
     return NodePath.get({
       parentPath: this.parentPath,
       parent: node,
diff --git a/packages/babel-traverse/src/path/modification.js b/packages/babel-traverse/src/path/modification.js
index 1402f1e3ae41..f425425efe4a 100644
--- a/packages/babel-traverse/src/path/modification.js
+++ b/packages/babel-traverse/src/path/modification.js
@@ -217,7 +217,7 @@ export function unshiftContainer(listKey, nodes) {
     container: this.node[listKey],
     listKey,
     key: 0,
-  });
+  }).setContext(this.context);
 
   return path._containerInsertBefore(nodes);
 }
@@ -237,7 +237,7 @@ export function pushContainer(listKey, nodes) {
     container: container,
     listKey,
     key: container.length,
-  });
+  }).setContext(this.context);
 
   return path.replaceWithMultiple(nodes);
 }
diff --git a/packages/babel-traverse/test/modification.js b/packages/babel-traverse/test/modification.js
index 615845547660..d26e4e68ad8f 100644
--- a/packages/babel-traverse/test/modification.js
+++ b/packages/babel-traverse/test/modification.js
@@ -49,6 +49,26 @@ describe("modification", function () {
         },
       });
     });
+
+    it("should set the correct path.context", function () {
+      expect.assertions(2);
+
+      const ast = parse("[b];");
+      traverse(ast, {
+        skipKeys: ["consequent"],
+        ExpressionStatement(path) {
+          path.traverse({ Identifier() {}, skipKeys: [] });
+
+          const arr = path.get("expression");
+          const x = arr.pushContainer("elements", [
+            { type: "Identifier", name: "x" },
+          ])[0];
+
+          expect(x.node.name).toBe("x");
+          expect(x.opts.skipKeys).toEqual(["consequent"]);
+        },
+      });
+    });
   });
   describe("unshiftContainer", function () {
     it("unshifts identifier into params", function () {
@@ -78,6 +98,26 @@ describe("modification", function () {
         },
       });
     });
+
+    it("should set the correct path.context", function () {
+      expect.assertions(2);
+
+      const ast = parse("[b];");
+      traverse(ast, {
+        skipKeys: ["consequent"],
+        ExpressionStatement(path) {
+          path.traverse({ Identifier() {}, skipKeys: [] });
+
+          const arr = path.get("expression");
+          const x = arr.unshiftContainer("elements", [
+            { type: "Identifier", name: "x" },
+          ])[0];
+
+          expect(x.node.name).toBe("x");
+          expect(x.opts.skipKeys).toEqual(["consequent"]);
+        },
+      });
+    });
   });
 
   describe("insertBefore", function () {