Fix flow comments plugin issues

packages/babel-plugin-transform-flow-comments/src/index.js

@@ -6,27 +6,73 @@ import generateCode from "@babel/generator";
 export default declare(api => {
   api.assertVersion(7);
 
-  function attachComment(path, comment) {
-    let attach = path.getPrevSibling();
-    let where = "trailing";
-    if (!attach.node) {
-      attach = path.parentPath;
+  function commentFromString(comment) {
+    return typeof comment === "string"
+      ? { type: "CommentBlock", value: comment }
+      : comment;
+  }
+
+  function attachComment({
+    ofPath,
+    toPath,
+    where = "trailing",
+    optional = false,
+    comments = generateComment(ofPath, optional),
+    keepType = false,
+  }) {
+    if (!toPath || !toPath.node) {
+      toPath = ofPath.getPrevSibling();
+      where = "trailing";
+    }
+    if (!toPath.node) {
+      toPath = ofPath.getNextSibling();
+      where = "leading";
+    }
+    if (!toPath.node) {
+      toPath = ofPath.parentPath;
       where = "inner";
     }
-    attach.addComment(where, comment);
-    path.remove();
+    if (!Array.isArray(comments)) {
+      comments = [comments];
+    }
+    comments = comments.map(commentFromString);
+    if (!keepType && ofPath && ofPath.node) {
+      // Removes the node at `ofPath` while conserving the comments attached
+      // to it.
+      const node = ofPath.node;
+      const parent = ofPath.parentPath;
+      const prev = ofPath.getPrevSibling();
+      const next = ofPath.getNextSibling();
+      const isSingleChild = !(prev.node || next.node);
+      const leading = node.leadingComments;
+      const trailing = node.trailingComments;
+
+      if (isSingleChild && leading) {
+        parent.addComments("inner", leading);
+      }
+      toPath.addComments(where, comments);
+      ofPath.remove();
+      if (isSingleChild && trailing) {
+        parent.addComments("inner", trailing);
+      }
+    } else {
+      toPath.addComments(where, comments);
+    }
   }
 
-  function wrapInFlowComment(path, parent) {
-    attachComment(path, generateComment(path, parent));
+  function wrapInFlowComment(path) {
+    attachComment({
+      ofPath: path,
+      comments: generateComment(path, path.parent.optional),
+    });
   }
 
-  function generateComment(path, parent) {
+  function generateComment(path, optional) {
     let comment = path
       .getSource()
       .replace(/\*-\//g, "*-ESCAPED/")
       .replace(/\*\//g, "*-/");
-    if (parent && parent.optional) comment = "?" + comment;
+    if (optional) comment = "?" + comment;
     if (comment[0] !== ":") comment = ":: " + comment;
     return comment;
   }
@@ -42,28 +88,32 @@ export default declare(api => {
     visitor: {
       TypeCastExpression(path) {
         const { node } = path;
-        path
-          .get("expression")
-          .addComment("trailing", generateComment(path.get("typeAnnotation")));
+        attachComment({
+          ofPath: path.get("typeAnnotation"),
+          toPath: path.get("expression"),
+          keepType: true,
+        });
         path.replaceWith(t.parenthesizedExpression(node.expression));
       },
 
       // support function a(b?) {}
       Identifier(path) {
-        if (path.parentPath.isFlow()) {
-          return;
-        }
-
+        if (path.parentPath.isFlow()) return;
         const { node } = path;
         if (node.typeAnnotation) {
-          const typeAnnotation = path.get("typeAnnotation");
-          path.addComment("trailing", generateComment(typeAnnotation, node));
-          typeAnnotation.remove();
+          attachComment({
+            ofPath: path.get("typeAnnotation"),
+            toPath: path,
+            optional: node.optional || node.typeAnnotation.optional,
+          });
           if (node.optional) {
             node.optional = false;
           }
         } else if (node.optional) {
-          path.addComment("trailing", ":: ?");
+          attachComment({
+            toPath: path,
+            comments: ":: ?",
+          });
           node.optional = false;
         }
       },
@@ -81,58 +131,51 @@ export default declare(api => {
       Function(path) {
         if (path.isDeclareFunction()) return;
         const { node } = path;
-        if (node.returnType) {
-          const returnType = path.get("returnType");
-          const typeAnnotation = returnType.get("typeAnnotation");
-          const block = path.get("body");
-          block.addComment(
-            "leading",
-            generateComment(returnType, typeAnnotation.node),
-          );
-          returnType.remove();
-        }
         if (node.typeParameters) {
-          const typeParameters = path.get("typeParameters");
-          const id = path.get("id");
-          id.addComment(
-            "trailing",
-            generateComment(typeParameters, typeParameters.node),
-          );
-          typeParameters.remove();
+          attachComment({
+            ofPath: path.get("typeParameters"),
+            toPath: path.get("id"),
+            optional: node.typeParameters.optional,
+          });
+        }
+        if (node.returnType) {
+          attachComment({
+            ofPath: path.get("returnType"),
+            toPath: path.get("body"),
+            where: "leading",
+            optional: node.returnType.typeAnnotation.optional,
+          });
         }
       },
 
       // support for `class X { foo: string }` - #4622
       ClassProperty(path) {
-        const { node, parent } = path;
+        const { node } = path;
         if (!node.value) {
-          wrapInFlowComment(path, parent);
+          wrapInFlowComment(path);
         } else if (node.typeAnnotation) {
-          const typeAnnotation = path.get("typeAnnotation");
-          path
-            .get("key")
-            .addComment(
-              "trailing",
-              generateComment(typeAnnotation, typeAnnotation.node),
-            );
-          typeAnnotation.remove();
+          attachComment({
+            ofPath: path.get("typeAnnotation"),
+            toPath: path.get("key"),
+            optional: node.typeAnnotation.optional,
+          });
         }
       },
 
       // support `export type a = {}` - #8 Error: You passed path.replaceWith() a falsy node
       ExportNamedDeclaration(path) {
-        const { node, parent } = path;
+        const { node } = path;
         if (node.exportKind !== "type" && !t.isFlow(node.declaration)) {
           return;
         }
-        wrapInFlowComment(path, parent);
+        wrapInFlowComment(path);
       },
 
       // support `import type A` and `import typeof A` #10
       ImportDeclaration(path) {
-        const { node, parent } = path;
+        const { node } = path;
         if (isTypeImport(node.importKind)) {
-          wrapInFlowComment(path, parent);
+          wrapInFlowComment(path);
           return;
         }
 
@@ -148,61 +191,88 @@ export default declare(api => {
         if (typeSpecifiers.length > 0) {
           const typeImportNode = t.cloneNode(node);
           typeImportNode.specifiers = typeSpecifiers;
+          const comment = `:: ${generateCode(typeImportNode).code}`;
 
           if (nonTypeSpecifiers.length > 0) {
-            path.addComment(
-              "trailing",
-              `:: ${generateCode(typeImportNode).code}`,
-            );
+            attachComment({ toPath: path, comments: comment });
           } else {
-            attachComment(path, `:: ${generateCode(typeImportNode).code}`);
+            attachComment({ ofPath: path, comments: comment });
           }
         }
       },
       ObjectPattern(path) {
         const { node } = path;
         if (node.typeAnnotation) {
-          const typeAnnotation = path.get("typeAnnotation");
-          path.addComment(
-            "trailing",
-            generateComment(typeAnnotation, typeAnnotation.node),
-          );
-          typeAnnotation.remove();
+          attachComment({
+            ofPath: path.get("typeAnnotation"),
+            toPath: path,
+            optional: node.optional || node.typeAnnotation.optional,
+          });
         }
       },
 
       Flow(path) {
-        const { parent } = path;
-        wrapInFlowComment(path, parent);
+        wrapInFlowComment(path);
       },
 
       Class(path) {
         const { node } = path;
-        if (node.typeParameters || node.implements) {
-          const comments = [];
-          if (node.typeParameters) {
-            const typeParameters = path.get("typeParameters");
+        let comments = [];
+        if (node.typeParameters) {
+          const typeParameters = path.get("typeParameters");
+          comments.push(
+            generateComment(typeParameters, node.typeParameters.optional),
+          );
+          const trailingComments = node.typeParameters.trailingComments;
+          if (trailingComments) {
+            comments.push(...trailingComments);
+          }
+          typeParameters.remove();
+        }
+
+        if (node.superClass) {
+          if (comments.length > 0) {
+            attachComment({
+              toPath: path.get("id"),
+              comments: comments,
+            });
+            comments = [];
+          }
+
+          if (node.superTypeParameters) {
+            const superTypeParameters = path.get("superTypeParameters");
             comments.push(
-              generateComment(typeParameters, typeParameters.node).replace(
-                /^:: /,
-                "",
+              generateComment(
+                superTypeParameters,
+                superTypeParameters.node.optional,
               ),
             );
-            typeParameters.remove();
+            superTypeParameters.remove();
           }
-          if (node.implements) {
-            const impls = path.get("implements");
-            comments.push(
-              "implements " +
-                impls
-                  .map(impl => generateComment(impl).replace(/^:: /, ""))
-                  .join(", "),
-            );
-            delete node["implements"];
+        }
+
+        if (node.implements) {
+          const impls = path.get("implements");
+          const comment =
+            "implements " +
+            impls
+              .map(impl => generateComment(impl).replace(/^:: /, ""))
+              .join(", ");
+          delete node["implements"];
+
+          if (comments.length === 1) {
+            comments[0] += ` ${comment}`;
+          } else {
+            comments.push(`:: ${comment}`);
           }
+        }
 
-          const block = path.get("body");
-          block.addComment("leading", ":: " + comments.join(" "));
+        if (comments.length > 0) {
+          attachComment({
+            toPath: path.get("body"),
+            where: "leading",
+            comments: comments,
+          });
         }
       },
     },

packages/babel-plugin-transform-flow-comments/test/fixtures/flow-comments/class-extends-type-parameter/input.js

@@ -0,0 +1 @@
+class Foo extends Bar<T> {}

packages/babel-plugin-transform-flow-comments/test/fixtures/flow-comments/class-extends-type-parameter/output.js

@@ -0,0 +1,3 @@
+class Foo extends Bar
+/*:: <T>*/
+{}

packages/babel-plugin-transform-flow-comments/test/fixtures/flow-comments/class-type-parameter-extends-type-parameter/input.js

@@ -0,0 +1,3 @@
+class Foo<T> /* inner */ extends Bar<R> {}
+
+/*a*/class /*b*/Baz/*c*/<T>/*d*/extends /*e*/Bar/*f*/<R>/*g*/ {/*h*/}/*i*/

packages/babel-plugin-transform-flow-comments/test/fixtures/flow-comments/class-type-parameter-extends-type-parameter/output.js

@@ -0,0 +1,30 @@
+class Foo
+/*:: <T>*/
+
+/* inner */
+extends Bar
+/*:: <R>*/
+{}
+/*a*/
+
+
+class
+/*b*/
+Baz
+/*c*/
+
+/*:: <T>*/
+
+/*d*/
+extends
+/*e*/
+Bar
+/*f*/
+
+/*:: <R>*/
+
+/*g*/
+{}
+/*h*/
+
+/*i*/

packages/babel-plugin-transform-flow-comments/test/fixtures/flow-comments/class-type-parameter-extends/input.js

@@ -0,0 +1 @@
+class Foo<T> extends Bar {}

packages/babel-plugin-transform-flow-comments/test/fixtures/flow-comments/class-type-parameter-extends/output.js

@@ -0,0 +1,3 @@
+class Foo
+/*:: <T>*/
+extends Bar {}

packages/babel-plugin-transform-flow-comments/test/fixtures/flow-comments/preserve-comments-order/input.js

@@ -0,0 +1,7 @@
+/*a*/
+type Foo = number;
+/*b*/
+var foo;
+/*c*/
+type Bar = number;
+/*d*/

packages/babel-plugin-transform-flow-comments/test/fixtures/flow-comments/preserve-comments-order/output.js

@@ -0,0 +1,11 @@
+/*a*/
+
+/*:: type Foo = number;*/
+
+/*b*/
+var foo;
+/*c*/
+
+/*:: type Bar = number;*/
+
+/*d*/