Add display name after create context

packages/babel-plugin-transform-react-display-name/src/index.js

@@ -2,29 +2,111 @@ import { declare } from "@babel/helper-plugin-utils";
 import path from "path";
 import { types as t } from "@babel/core";
 
-export default declare(api => {
-  api.assertVersion(7);
-
-  function addDisplayName(id, call) {
-    const props = call.arguments[0].properties;
-    let safe = true;
-
-    for (let i = 0; i < props.length; i++) {
-      const prop = props[i];
-      const key = t.toComputedKey(prop);
-      if (t.isLiteral(key, { value: "displayName" })) {
-        safe = false;
-        break;
-      }
+function addDisplayNameInCreateClass(id, call) {
+  const props = call.arguments[0].properties;
+  let safe = true;
+
+  for (let i = 0; i < props.length; i++) {
+    const prop = props[i];
+    const key = t.toComputedKey(prop);
+    if (t.isLiteral(key, { value: "displayName" })) {
+      safe = false;
+      break;
     }
+  }
 
-    if (safe) {
-      props.unshift(
-        t.objectProperty(t.identifier("displayName"), t.stringLiteral(id)),
-      );
+  if (safe) {
+    props.unshift(
+      t.objectProperty(t.identifier("displayName"), t.stringLiteral(id)),
+    );
+  }
+}
+
+function getDisplayNameReferenceIdentifier(
+  path: NodePath<t.CallExpression>,
+): ?t.Node {
+  let id;
+
+  // crawl up the ancestry looking for possible candidates for displayName inference
+  path.find(function (path) {
+    if (path.isAssignmentExpression()) {
+      id = path.node.left;
+    } else if (path.isObjectProperty()) {
+      id = path.node.key;
+    } else if (path.isVariableDeclarator()) {
+      id = path.node.id;
+    } else if (path.isStatement()) {
+      // we've hit a statement, we should stop crawling up
+      return true;
     }
+
+    // we've got an id! no need to continue
+    if (id) return true;
+  });
+
+  // ensure that we have an identifier we can inherit from
+  if (!id) return;
+
+  // foo.bar -> bar
+  if (t.isMemberExpression(id)) {
+    id = id.property;
   }
 
+  return id;
+}
+
+function isCreateContext(node) {
+  let callee;
+  return (
+    t.isCallExpression(node) &&
+    t.isMemberExpression((callee = node.callee)) &&
+    t.isIdentifier(callee.object, { name: "React" }) &&
+    (t.isIdentifier(callee.property, { name: "createContext" }) ||
+      t.isStringLiteral(callee.property, { value: "createContext" }))
+  );
+}
+
+function buildDisplayNameAssignment(ref, displayName) {
+  return t.assignmentExpression(
+    "=",
+    t.memberExpression(t.cloneNode(ref), t.identifier("displayName")),
+    t.stringLiteral(displayName),
+  );
+}
+
+function addDisplayNameAfterCreateContext(
+  id,
+  path: t.NodePath<t.CallExpression>,
+) {
+  const { parentPath } = path;
+  if (parentPath.isVariableDeclarator()) {
+    // FooContext = React.createContext()
+    const ref = parentPath.node.id;
+    // parentPath.parentPath must be a VariableDeclaration because getDisplayNameReferenceIdentifier
+    // does not support patterns
+    parentPath.parentPath.insertAfter(buildDisplayNameAssignment(ref, id));
+  } else if (parentPath.isAssignmentExpression()) {
+    // var FooContext = React.createContext()
+    const ref = parentPath.node.left;
+    parentPath.insertAfter(buildDisplayNameAssignment(ref, id));
+  } else {
+    // (ref = React.createContext(), ref.displayName = "id", ref)
+    const ref = path.scope.generateUidIdentifier("ref");
+    path.replaceWith(
+      t.sequenceExpression([
+        t.assignmentExpression("=", ref, path.node),
+        buildDisplayNameAssignment(ref, id),
+        t.cloneNode(ref),
+      ]),
+    );
+  }
+}
+
+const createContextVisited = new WeakSet();
+
+export default declare(api => {
+  api.assertVersion(7);
+
   const isCreateClassCallExpression =
     t.buildMatchMemberExpression("React.createClass");
   const isCreateClassAddon = callee => callee.name === "createReactClass";
@@ -66,44 +148,28 @@ export default declare(api => {
             displayName = path.basename(path.dirname(filename));
           }
 
-          addDisplayName(displayName, node.declaration);
+          addDisplayNameInCreateClass(displayName, node.declaration);
         }
       },
 
       CallExpression(path) {
         const { node } = path;
-        if (!isCreateClass(node)) return;
-
-        let id;
-
-        // crawl up the ancestry looking for possible candidates for displayName inference
-        path.find(function (path) {
-          if (path.isAssignmentExpression()) {
-            id = path.node.left;
-          } else if (path.isObjectProperty()) {
-            id = path.node.key;
-          } else if (path.isVariableDeclarator()) {
-            id = path.node.id;
-          } else if (path.isStatement()) {
-            // we've hit a statement, we should stop crawling up
-            return true;
+        if (isCreateClass(node)) {
+          const id = getDisplayNameReferenceIdentifier(path);
+          // identifiers are the only thing we can reliably get a name from
+          if (t.isIdentifier(id)) {
+            addDisplayNameInCreateClass(id.name, node);
+          }
+        } else if (isCreateContext(node)) {
+          if (createContextVisited.has(node)) {
+            return;
+          }
+          createContextVisited.add(node);
+          const id = getDisplayNameReferenceIdentifier(path);
+          // identifiers are the only thing we can reliably get a name from
+          if (t.isIdentifier(id)) {
+            addDisplayNameAfterCreateContext(id.name, path);
           }
-
-          // we've got an id! no need to continue
-          if (id) return true;
-        });
-
-        // ensure that we have an identifier we can inherit from
-        if (!id) return;
-
-        // foo.bar -> bar
-        if (t.isMemberExpression(id)) {
-          id = id.property;
-        }
-
-        // identifiers are the only thing we can reliably get a name from
-        if (t.isIdentifier(id)) {
-          addDisplayName(id.name, node);
         }
       },
     },

packages/babel-plugin-transform-react-display-name/test/fixtures/display-name-context/assignment-expression-and-display-name/input.js

@@ -0,0 +1,2 @@
+ThemeContext = React.createContext("light");
+ThemeContext.displayName = "CustomThemeContext";

packages/babel-plugin-transform-react-display-name/test/fixtures/display-name-context/assignment-expression-and-display-name/output.js

@@ -0,0 +1,3 @@
+ThemeContext = React.createContext("light");
+ThemeContext.displayName = "ThemeContext";
+ThemeContext.displayName = "CustomThemeContext";

packages/babel-plugin-transform-react-display-name/test/fixtures/display-name-context/assignment-expression/input.js

@@ -0,0 +1 @@
+ThemeContext = React.createContext("light");

packages/babel-plugin-transform-react-display-name/test/fixtures/display-name-context/assignment-expression/output.js

@@ -0,0 +1,2 @@
+ThemeContext = React.createContext("light");
+ThemeContext.displayName = "ThemeContext";

packages/babel-plugin-transform-react-display-name/test/fixtures/display-name-context/nested/input.js

@@ -0,0 +1 @@
+var enhancedContext = qux(React.createContext("light"));

packages/babel-plugin-transform-react-display-name/test/fixtures/display-name-context/nested/output.js

@@ -0,0 +1 @@
+var enhancedContext = qux((_ref = React.createContext("light"), _ref.displayName = "enhancedContext", _ref));

packages/babel-plugin-transform-react-display-name/test/fixtures/display-name-context/object-property/input.js

@@ -0,0 +1,3 @@
+({
+  ThemeContext: React.createContext("light")
+});

packages/babel-plugin-transform-react-display-name/test/fixtures/display-name-context/object-property/output.js

@@ -0,0 +1,3 @@
+({
+  ThemeContext: (_ref = React.createContext("light"), _ref.displayName = "ThemeContext", _ref)
+});

packages/babel-plugin-transform-react-display-name/test/fixtures/display-name-context/options.json

@@ -0,0 +1,3 @@
+{
+  "plugins": ["transform-react-display-name"]
+}

packages/babel-plugin-transform-react-display-name/test/fixtures/display-name-context/variable-declarator/input.js

@@ -0,0 +1 @@
+var ThemeContext = React.createContext("light");

packages/babel-plugin-transform-react-display-name/test/fixtures/display-name-context/variable-declarator/output.js

@@ -0,0 +1,2 @@
+var ThemeContext = React.createContext("light");
+ThemeContext.displayName = "ThemeContext"