From 7bf1803a11fef7d12d2827bd9cf619125a1f6f0b Mon Sep 17 00:00:00 2001
From: fisker Cheung <lionkay@gmail.com>
Date: Fri, 21 Aug 2020 09:10:14 +0800
Subject: [PATCH] JS: DRY `hasPragma` & `insertPragma` (#9002)

---
 src/language-js/pragma.js                     | 45 ++++++++++--------
 .../js/__snapshots__/jsfmt.spec.js.snap       | 47 +++++++++++++++++++
 tests/misc/insert-pragma/js/empty.js          |  0
 .../js/shebang-only-with-pragma.js            |  2 +
 .../js/shebang-only-without-pragma.js         |  1 +
 5 files changed, 75 insertions(+), 20 deletions(-)
 create mode 100644 tests/misc/insert-pragma/js/empty.js
 create mode 100644 tests/misc/insert-pragma/js/shebang-only-with-pragma.js
 create mode 100644 tests/misc/insert-pragma/js/shebang-only-without-pragma.js

diff --git a/src/language-js/pragma.js b/src/language-js/pragma.js
index bd1215872..4bb08133f 100644
--- a/src/language-js/pragma.js
+++ b/src/language-js/pragma.js
@@ -1,36 +1,41 @@
 "use strict";
 
-const docblock = require("jest-docblock");
+const { parseWithComments, strip, extract, print } = require("jest-docblock");
 const { getShebang } = require("../common/util");
 
-function hasPragma(text) {
+function parseDocBlock(text) {
   const shebang = getShebang(text);
   if (shebang) {
     text = text.slice(shebang.length + 1);
   }
-  const pragmas = Object.keys(docblock.parse(docblock.extract(text)));
+
+  const docBlock = extract(text);
+  const { pragmas, comments } = parseWithComments(docBlock);
+
+  return { shebang, text, pragmas, comments };
+}
+
+function hasPragma(text) {
+  const pragmas = Object.keys(parseDocBlock(text).pragmas);
   return pragmas.includes("prettier") || pragmas.includes("format");
 }
 
-function insertPragma(text) {
-  const shebang = getShebang(text);
-  if (shebang) {
-    text = text.slice(shebang.length + 1);
-  }
-  const parsedDocblock = docblock.parseWithComments(docblock.extract(text));
-  const pragmas = { format: "", ...parsedDocblock.pragmas };
-  const newDocblock = docblock
-    .print({
-      pragmas,
-      comments: parsedDocblock.comments.replace(/^(\s+?\r?\n)+/, ""), // remove leading newlines
-    })
-    .replace(/(\r\n|\r)/g, "\n"); // normalise newlines (mitigate use of os.EOL by jest-docblock)
-  const strippedText = docblock.strip(text);
-  const separatingNewlines = strippedText.startsWith("\n") ? "\n" : "\n\n";
+function insertPragma(originalText) {
+  const { shebang, text, pragmas, comments } = parseDocBlock(originalText);
+  const strippedText = strip(text);
+
+  const docBlock = print({
+    pragmas: {
+      format: "",
+      ...pragmas,
+    },
+    comments: comments.trimStart(),
+  }).replace(/\r\n?/g, "\n"); // normalise newlines (mitigate use of os.EOL by jest-docblock)
+
   return (
     (shebang ? `${shebang}\n` : "") +
-    newDocblock +
-    separatingNewlines +
+    docBlock +
+    (strippedText.startsWith("\n") ? "\n" : "\n\n") +
     strippedText
   );
 }
diff --git a/tests/misc/insert-pragma/js/__snapshots__/jsfmt.spec.js.snap b/tests/misc/insert-pragma/js/__snapshots__/jsfmt.spec.js.snap
index 70f15fbc8..8ec5cc9b2 100644
--- a/tests/misc/insert-pragma/js/__snapshots__/jsfmt.spec.js.snap
+++ b/tests/misc/insert-pragma/js/__snapshots__/jsfmt.spec.js.snap
@@ -1,5 +1,19 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
+exports[`empty.js - {"insertPragma":true} format 1`] = `
+====================================options=====================================
+insertPragma: true
+parsers: ["flow", "babel", "typescript"]
+printWidth: 80
+                                                                                | printWidth
+=====================================input======================================
+
+=====================================output=====================================
+/** @format */
+
+================================================================================
+`;
+
 exports[`module-with-late-docblock.js - {"insertPragma":true} format 1`] = `
 ====================================options=====================================
 insertPragma: true
@@ -139,6 +153,39 @@ function foo(bar) {
 ================================================================================
 `;
 
+exports[`shebang-only-with-pragma.js - {"insertPragma":true} format 1`] = `
+====================================options=====================================
+insertPragma: true
+parsers: ["flow", "babel", "typescript"]
+printWidth: 80
+                                                                                | printWidth
+=====================================input======================================
+#!/usr/bin/env node
+/** @format */
+
+=====================================output=====================================
+#!/usr/bin/env node
+/** @format */
+
+================================================================================
+`;
+
+exports[`shebang-only-without-pragma.js - {"insertPragma":true} format 1`] = `
+====================================options=====================================
+insertPragma: true
+parsers: ["flow", "babel", "typescript"]
+printWidth: 80
+                                                                                | printWidth
+=====================================input======================================
+#!/usr/bin/env node
+
+=====================================output=====================================
+#!/usr/bin/env node
+/** @format */
+
+================================================================================
+`;
+
 exports[`shebang-with-pragma.js - {"insertPragma":true} format 1`] = `
 ====================================options=====================================
 insertPragma: true
diff --git a/tests/misc/insert-pragma/js/empty.js b/tests/misc/insert-pragma/js/empty.js
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/misc/insert-pragma/js/shebang-only-with-pragma.js b/tests/misc/insert-pragma/js/shebang-only-with-pragma.js
new file mode 100644
index 000000000..507b6a49e
--- /dev/null
+++ b/tests/misc/insert-pragma/js/shebang-only-with-pragma.js
@@ -0,0 +1,2 @@
+#!/usr/bin/env node
+/** @format */
diff --git a/tests/misc/insert-pragma/js/shebang-only-without-pragma.js b/tests/misc/insert-pragma/js/shebang-only-without-pragma.js
new file mode 100644
index 000000000..908ba8417
--- /dev/null
+++ b/tests/misc/insert-pragma/js/shebang-only-without-pragma.js
@@ -0,0 +1 @@
+#!/usr/bin/env node