Fix Missing Comments (#24)

* Keep comments by default when using generateFlatAST * Add comment related fields * Keep comments when replacing and deleting nodes * 1.5.2
PerimeterX · Sep 10, 2023 · fc64a65 · fc64a65
1 parent 010418d
commit fc64a65
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 5 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "flast",
-  "version": "1.5.1",
+  "version": "1.5.2",
   "description": "Flatten JS AST",
   "main": "src/index.js",
   "scripts": {

diff --git a/src/arborist.js b/src/arborist.js
@@ -82,7 +82,11 @@ const Arborist = class {
 				if (rootNodeReplacement) {
 					++changesCounter;
 					this.log(`[+] Applying changes to the root node...`);
+					const leadingComments =  rootNode.leadingComments || [];
+					const trailingComments = rootNode.trailingComments || [];
 					rootNode = rootNodeReplacement[1];
+					if (leadingComments.length) rootNode.leadingComments = (rootNode.leadingComments || []).concat(leadingComments);
+					if (trailingComments.length) rootNode.trailingComments = (rootNode.trailingComments || []).concat(trailingComments);
 				} else {
 					for (const targetNodeId of this.markedForDeletion) {
 						try {
@@ -91,11 +95,18 @@ const Arborist = class {
 								const parent = targetNode.parentNode;
 								if (parent[targetNode.parentKey] === targetNode) {
 									parent[targetNode.parentKey] = undefined;
+									const comments = (targetNode.leadingComments || []).concat(targetNode.trailingComments || []);
+									if (comments.length) parent.trailingComments = (parent.trailingComments || []).concat(comments);
 									++changesCounter;
 								} else if (Array.isArray(parent[targetNode.parentKey])) {
 									const idx = parent[targetNode.parentKey].indexOf(targetNode);
 									parent[targetNode.parentKey][idx] = undefined;
 									parent[targetNode.parentKey] = parent[targetNode.parentKey].filter(n => n);
+									const comments = (targetNode.leadingComments || []).concat(targetNode.trailingComments || []);
+									if (comments.length) {
+										const targetParent = idx > 0 ? parent[targetNode.parentKey][idx - 1] : parent[targetNode.parentKey].length > 1 ? parent[targetNode.parentKey][idx + 1] : parent;
+										targetParent.trailingComments = (targetParent.trailingComments || []).concat(comments);
+									}
 									++changesCounter;
 								}
 							}
@@ -109,10 +120,22 @@ const Arborist = class {
 								const parent = targetNode.parentNode;
 								if (parent[targetNode.parentKey] === targetNode) {
 									parent[targetNode.parentKey] = replacementNode;
+									const leadingComments =  targetNode.leadingComments || [];
+									const trailingComments = targetNode.trailingComments || [];
+									if (leadingComments.length) replacementNode.leadingComments = (replacementNode.leadingComments || []).concat(leadingComments);
+									if (trailingComments.length) replacementNode.trailingComments = (replacementNode.trailingComments || []).concat(trailingComments);
 									++changesCounter;
 								} else if (Array.isArray(parent[targetNode.parentKey])) {
 									const idx = parent[targetNode.parentKey].indexOf(targetNode);
 									parent[targetNode.parentKey][idx] = replacementNode;
+									const comments = (targetNode.leadingComments || []).concat(targetNode.trailingComments || []);
+									if (idx > 0) {
+										const commentsTarget = parent[targetNode.parentKey][idx - 1];
+										commentsTarget.trailingComments = (commentsTarget.trailingComments || []).concat(comments);
+									} else if (parent[targetNode.parentKey].length > 1) {
+										const commentsTarget = parent[targetNode.parentKey][idx + 1];
+										commentsTarget.leadingComments = (commentsTarget.leadingComments || []).concat(comments);
+									} else parent.trailingComments = (parent.trailingComments || []).concat(comments);
 									++changesCounter;
 								}
 							}

diff --git a/src/flast.js b/src/flast.js
@@ -1,5 +1,5 @@
 const {parse} = require('espree');
-const {generate} = require('escodegen');
+const {generate, attachComments} = require('escodegen');
 const estraverse = require('estraverse');
 const {analyze} = require('eslint-scope');
 
@@ -12,7 +12,9 @@ const sourceType = 'module';
  * @return {ASTNode} The root of the AST
  */
 function parseCode(inputCode, opts = {}) {
-	return parse(inputCode, {ecmaVersion, comment: true, range: true, ...opts});
+	const rootNode = parse(inputCode, {ecmaVersion, comment: true, range: true, ...opts});
+	if (rootNode.tokens) attachComments(rootNode, rootNode.comments, rootNode.tokens);
+	return rootNode;
 }
 
 const excludedParentKeys = [
@@ -51,6 +53,8 @@ const generateFlatASTDefaultOptions = {
 	// Options for the espree parser
 	parseOpts: {
 		sourceType,
+		comment: true,
+		tokens: true,
 	},
 };
 

diff --git a/src/types.js b/src/types.js
@@ -11,6 +11,7 @@ const {Scope} = require('eslint-scope');
  * @property {ASTNode} [callee]
  * @property {ASTNode[]} [cases]
  * @property {ASTNode[]} [childNodes]
+ * @property {Object[]} [comments]
  * @property {boolean} [computed]
  * @property {ASTNode} [consequent]
  * @property {string} [cooked]
@@ -35,6 +36,7 @@ const {Scope} = require('eslint-scope');
  * @property {ASTNode} [key]
  * @property {string} [kind]
  * @property {ASTNode} [label]
+ * @property {Object[]} [leadingComments]
  * @property {ASTNode} [left]
  * @property {number[]} [lineage] The nodeIds of all parent nodes
  * @property {ASTNode} [local]
@@ -68,6 +70,8 @@ const {Scope} = require('eslint-scope');
  * @property {ASTNode} [superClass]
  * @property {boolean} [tail]
  * @property {ASTNode} [test]
+ * @property {ASTNode} [tokens]
+ * @property {Object[]} [trailingComments]
  * @property {ASTNode} [update]
  * @property {ASTNode|string|number|boolean} [value]
  */