keep script directives at top of file, fixes #185

src/preprocessors/preprocessor.ts

@@ -1,8 +1,8 @@
 import { ParserOptions, parse as babelParser } from '@babel/parser';
-import traverse, { NodePath } from '@babel/traverse';
-import { ImportDeclaration, isTSModuleDeclaration } from '@babel/types';
+import { Directive, ImportDeclaration } from '@babel/types';
 
 import { PrettierOptions } from '../types';
+import { extractASTNodes } from '../utils/extract-ast-nodes';
 import { getCodeFromAst } from '../utils/get-code-from-ast';
 import { getExperimentalParserPlugins } from '../utils/get-experimental-parser-plugins';
 import { getSortedNodes } from '../utils/get-sorted-nodes';
@@ -17,7 +17,6 @@ export function preprocessor(code: string, options: PrettierOptions) {
         importOrderSortSpecifiers,
     } = options;
 
-    const importNodes: ImportDeclaration[] = [];
     const parserOptions: ParserOptions = {
         sourceType: 'module',
         plugins: getExperimentalParserPlugins(importOrderParserPlugins),
@@ -26,16 +25,11 @@ export function preprocessor(code: string, options: PrettierOptions) {
     const ast = babelParser(code, parserOptions);
     const interpreter = ast.program.interpreter;
 
-    traverse(ast, {
-        ImportDeclaration(path: NodePath<ImportDeclaration>) {
-            const tsModuleParent = path.findParent((p) =>
-                isTSModuleDeclaration(p),
-            );
-            if (!tsModuleParent) {
-                importNodes.push(path.node);
-            }
-        },
-    });
+    const {
+        importNodes,
+        directives,
+    }: { importNodes: ImportDeclaration[]; directives: Directive[] } =
+        extractASTNodes(ast);
 
     // short-circuit if there are no import declaration
     if (importNodes.length === 0) return code;
@@ -48,5 +42,5 @@ export function preprocessor(code: string, options: PrettierOptions) {
         importOrderSortSpecifiers,
     });
 
-    return getCodeFromAst(allImports, code, interpreter);
+    return getCodeFromAst(allImports, directives, code, interpreter);
 }

src/utils/__tests__/get-code-from-ast.spec.ts

@@ -1,6 +1,10 @@
+import { parse as babelParser } from '@babel/core';
+import { ParserOptions } from '@babel/parser';
 import { format } from 'prettier';
 
+import { extractASTNodes } from '../extract-ast-nodes';
 import { getCodeFromAst } from '../get-code-from-ast';
+import { getExperimentalParserPlugins } from '../get-experimental-parser-plugins';
 import { getImportNodes } from '../get-import-nodes';
 import { getSortedNodes } from '../get-sorted-nodes';
 
@@ -22,7 +26,7 @@ import a from 'a';
         importOrderGroupNamespaceSpecifiers: false,
         importOrderSortSpecifiers: false,
     });
-    const formatted = getCodeFromAst(sortedNodes, code, null);
+    const formatted = getCodeFromAst(sortedNodes, [], code, null);
     expect(format(formatted, { parser: 'babel' })).toEqual(
         `// first comment
 // second comment
@@ -35,3 +39,29 @@ import z from "z";
 `,
     );
 });
+
+test('it renders directives correctly', () => {
+    const code = `
+    "use client";
+// first comment
+import b from 'b';
+import a from 'a';`;
+
+    const parserOptions: ParserOptions = {
+        sourceType: 'module',
+        plugins: getExperimentalParserPlugins([]),
+    };
+    const ast = babelParser(code, parserOptions);
+    if (!ast) throw new Error('ast is null');
+    const { directives, importNodes } = extractASTNodes(ast);
+
+    const formatted = getCodeFromAst(importNodes, directives, code, null);
+    expect(format(formatted, { parser: 'babel' })).toEqual(
+        `"use client";
+
+// first comment
+import b from "b";
+import a from "a";
+`,
+    );
+});

src/utils/extract-ast-nodes.ts

@@ -0,0 +1,31 @@
+import { ParseResult } from '@babel/parser';
+import traverse, { NodePath } from '@babel/traverse';
+import {
+    Directive,
+    File,
+    ImportDeclaration,
+    isTSModuleDeclaration,
+} from '@babel/types';
+
+export function extractASTNodes(ast: ParseResult<File>) {
+    const importNodes: ImportDeclaration[] = [];
+    const directives: Directive[] = [];
+    traverse(ast, {
+        Directive({ node }) {
+            directives.push(node);
+
+            // Trailing comments probably shouldn't be attached to the directive
+            node.trailingComments = null;
+        },
+
+        ImportDeclaration(path: NodePath<ImportDeclaration>) {
+            const tsModuleParent = path.findParent((p) =>
+                isTSModuleDeclaration(p),
+            );
+            if (!tsModuleParent) {
+                importNodes.push(path.node);
+            }
+        },
+    });
+    return { importNodes, directives };
+}

src/utils/get-code-from-ast.ts

@@ -1,5 +1,5 @@
 import generate from '@babel/generator';
-import { InterpreterDirective, Statement, file } from '@babel/types';
+import { Directive, InterpreterDirective, Statement, file } from '@babel/types';
 
 import { newLineCharacters } from '../constants';
 import { getAllCommentsFromNodes } from './get-all-comments-from-nodes';
@@ -12,12 +12,14 @@ import { removeNodesFromOriginalCode } from './remove-nodes-from-original-code';
  */
 export const getCodeFromAst = (
     nodes: Statement[],
+    directives: Directive[],
     originalCode: string,
     interpreter?: InterpreterDirective | null,
 ) => {
     const allCommentsFromImports = getAllCommentsFromNodes(nodes);
 
     const nodesToRemoveFromCode = [
+        ...directives,
         ...nodes,
         ...allCommentsFromImports,
         ...(interpreter ? [interpreter] : []),
@@ -31,7 +33,7 @@ export const getCodeFromAst = (
     const newAST = file({
         type: 'Program',
         body: nodes,
-        directives: [],
+        directives,
         sourceType: 'module',
         interpreter: interpreter,
         sourceFile: '',

src/utils/remove-nodes-from-original-code.ts

@@ -1,6 +1,7 @@
 import {
     CommentBlock,
     CommentLine,
+    Directive,
     ImportDeclaration,
     InterpreterDirective,
     Statement,
@@ -21,6 +22,7 @@ export const removeNodesFromOriginalCode = (
     nodes: (
         | Statement
         | CommentBlock
+        | Directive
         | CommentLine
         | ImportDeclaration
         | InterpreterDirective

tests/ImportsNotSeparated/__snapshots__/ppsi.spec.js.snap

@@ -151,6 +151,29 @@ function add(a: number, b: number) {
 
 `;
 
+exports[`imports-with-directives.ts - typescript-verify: imports-with-directives.ts 1`] = `
+'use strict';
+'use client';
+import otherthing from "@core/otherthing";
+import abc from "@core/abc";
+// Comment
+function add(a:number,b:number) {
+    return a + b;
+}
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+"use strict";
+"use client";
+
+import abc from "@core/abc";
+import otherthing from "@core/otherthing";
+
+// Comment
+function add(a: number, b: number) {
+    return a + b;
+}
+
+`;
+
 exports[`imports-with-file-level-comments.ts - typescript-verify: imports-with-file-level-comments.ts 1`] = `
 //@ts-ignore
 // I am file top level comments

tests/ImportsNotSeparated/imports-with-directives.ts

@@ -0,0 +1,8 @@
+'use strict';
+'use client';
+import otherthing from "@core/otherthing";
+import abc from "@core/abc";
+// Comment
+function add(a:number,b:number) {
+    return a + b;
+}

tests/ImportsSeparated/__snapshots__/ppsi.spec.js.snap

@@ -160,6 +160,35 @@ function add(a: number, b: number) {
 
 `;
 
+exports[`imports-with-directives.ts - typescript-verify: imports-with-directives.ts 1`] = `
+'use strict';
+'use client';
+
+// comment after directives
+import otherthing from "@core/otherthing";
+import abc from "@core/abc";
+
+// Comment
+
+function add(a:number,b:number) {
+  return a + b;
+}
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+"use strict";
+"use client";
+
+// comment after directives
+import abc from "@core/abc";
+import otherthing from "@core/otherthing";
+
+// Comment
+
+function add(a: number, b: number) {
+    return a + b;
+}
+
+`;
+
 exports[`imports-with-file-level-comments.ts - typescript-verify: imports-with-file-level-comments.ts 1`] = `
 //@ts-ignore
 // I am file top level comments

tests/ImportsSeparated/imports-with-directives.ts

@@ -0,0 +1,12 @@
+'use strict';
+'use client';
+
+// comment after directives
+import otherthing from "@core/otherthing";
+import abc from "@core/abc";
+
+// Comment
+
+function add(a:number,b:number) {
+  return a + b;
+}