Skip to content

Commit 463f2d1

Browse files
knutwannhedenknutwannheden
committed
Polish JS RemoveImport recipe
Now better handles trailing comments and comments on subsequent statements.
1 parent 3571561 commit 463f2d1

2 files changed

Lines changed: 185 additions & 26 deletions

File tree

‎rewrite-javascript/rewrite/src/javascript/remove-import.ts‎

Lines changed: 66 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -110,26 +110,23 @@ export class RemoveImport<P> extends JavaScriptVisitor<P> {
110110

111111
// Now process imports with knowledge of what's used
112112
return this.produceJavaScript<JS.CompilationUnit>(compilationUnit, p, async draft => {
113-
let nextStatementPrefix: J.Space | undefined;
113+
let previousWasRemoved = false;
114+
let hasKeptStatement = false;
114115

115-
draft.statements = await mapAsync(compilationUnit.statements, async (stmt, index) => {
116+
draft.statements = await mapAsync(compilationUnit.statements, async (stmt) => {
116117
const statement = stmt.element;
117118

118119
// Handle ES6 imports
119120
if (statement?.kind === JS.Kind.Import) {
120121
const jsImport = statement as JS.Import;
121122
const result = await this.processImport(jsImport, usedIdentifiers, usedTypes, p);
122123
if (result === undefined) {
123-
// Store the prefix for the next statement to avoid leaving blank lines
124-
if (index < compilationUnit.statements.length - 1) {
125-
const nextStmt = compilationUnit.statements[index + 1];
126-
if (nextStmt?.element) {
127-
nextStatementPrefix = jsImport.prefix;
128-
}
129-
}
130-
// Remove the entire import
124+
// Mark that we removed an import
125+
previousWasRemoved = true;
131126
return undefined;
132127
}
128+
previousWasRemoved = false;
129+
hasKeptStatement = true;
133130
return {...stmt, element: result};
134131
}
135132

@@ -139,29 +136,25 @@ export class RemoveImport<P> extends JavaScriptVisitor<P> {
139136
const varDecl = statement as J.VariableDeclarations;
140137
const result = await this.processRequireFromVarDecls(varDecl, usedIdentifiers, p);
141138
if (result === undefined) {
142-
// Store the prefix for the next statement to avoid leaving blank lines
143-
if (index < compilationUnit.statements.length - 1) {
144-
const nextStmt = compilationUnit.statements[index + 1];
145-
if (nextStmt?.element) {
146-
nextStatementPrefix = varDecl.prefix;
147-
}
148-
}
149-
// Remove the entire require statement
139+
// Mark that we removed a require
140+
previousWasRemoved = true;
150141
return undefined;
151142
}
143+
previousWasRemoved = false;
144+
hasKeptStatement = true;
152145
return {...stmt, element: result};
153146
}
154147

155-
// Apply stored prefix if this statement follows a removed import
156-
if (nextStatementPrefix && statement) {
157-
const updatedStatement = await this.visit(statement, p) as J;
158-
if (updatedStatement) {
159-
(updatedStatement as any).prefix = nextStatementPrefix;
160-
nextStatementPrefix = undefined;
161-
return {...stmt, element: updatedStatement};
162-
}
148+
// If the previous statement was removed, adjust this statement's prefix
149+
if (previousWasRemoved && statement) {
150+
previousWasRemoved = false;
151+
const updatedStatement = this.adjustPrefixAfterRemoval(statement, hasKeptStatement);
152+
hasKeptStatement = true;
153+
return {...stmt, element: updatedStatement};
163154
}
164155

156+
previousWasRemoved = false;
157+
hasKeptStatement = true;
165158
return stmt;
166159
});
167160

@@ -171,6 +164,53 @@ export class RemoveImport<P> extends JavaScriptVisitor<P> {
171164
});
172165
}
173166

167+
/**
168+
* Adjusts the prefix of a statement that follows a removed import/require.
169+
* Removes the import's line while preserving comments that belong to subsequent lines.
170+
*
171+
* If the prefix whitespace doesn't contain a newline, the first comment was inline
172+
* on the removed import line - we remove that comment and one line terminator from its suffix.
173+
*
174+
* @param statement The statement following the removed import
175+
* @param hasKeptPreviousStatement Whether there's a previous statement that wasn't removed
176+
*/
177+
private adjustPrefixAfterRemoval(statement: J, hasKeptPreviousStatement: boolean): J {
178+
const prefix = (statement as any).prefix;
179+
if (!prefix) {
180+
return statement;
181+
}
182+
183+
let whitespace = prefix.whitespace || '';
184+
let comments = prefix.comments || [];
185+
186+
// If the whitespace before the first comment doesn't contain a line terminator,
187+
// the first comment was inline on the same line as the removed import
188+
if (!/[\r\n]/.test(whitespace) && comments.length > 0) {
189+
// Remove the first comment (which was inline with the import)
190+
comments = comments.slice(1);
191+
// Clear the whitespace (it was just spacing before the inline comment)
192+
whitespace = '';
193+
} else if (!hasKeptPreviousStatement) {
194+
// At the beginning of file - remove all leading newlines
195+
// (The removed import was the first statement, so we don't want leading blank lines)
196+
whitespace = whitespace.replace(/^[\r\n]+/, '');
197+
}
198+
// else: There's a previous kept statement and no inline comment - keep whitespace as-is
199+
// (This preserves intentional blank lines between statements)
200+
201+
// Create a new Space with adjusted whitespace and comments
202+
const newPrefix = {
203+
...prefix,
204+
whitespace: whitespace,
205+
comments: comments
206+
};
207+
208+
return {
209+
...statement,
210+
prefix: newPrefix
211+
} as J;
212+
}
213+
174214
private async processImport(
175215
jsImport: JS.Import,
176216
usedIdentifiers: Set<string>,

‎rewrite-javascript/rewrite/test/javascript/remove-import.test.ts‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -413,4 +413,123 @@ describe('RemoveImport visitor', () => {
413413
);
414414
});
415415
});
416+
417+
describe('comment preservation', () => {
418+
test('should preserve comments on subsequent lines when removing import', async () => {
419+
const spec = new RecipeSpec();
420+
spec.recipe = fromVisitor(new RemoveImport("fs"));
421+
422+
//language=typescript
423+
await spec.rewriteRun(
424+
typescript(
425+
`
426+
import fs from 'fs';
427+
428+
// This comment belongs to the class below
429+
class MyClass {
430+
doSomething() {
431+
return 'test';
432+
}
433+
}
434+
`,
435+
`
436+
// This comment belongs to the class below
437+
class MyClass {
438+
doSomething() {
439+
return 'test';
440+
}
441+
}
442+
`
443+
)
444+
);
445+
});
446+
447+
test('should preserve multiple comments after removed import', async () => {
448+
const spec = new RecipeSpec();
449+
spec.recipe = fromVisitor(new RemoveImport("fs", "readFile"));
450+
451+
//language=typescript
452+
await spec.rewriteRun(
453+
typescript(
454+
`
455+
import {readFile} from 'fs';
456+
457+
/**
458+
* Documentation for the function
459+
* @returns {string} A test value
460+
*/
461+
function example() {
462+
return 'test';
463+
}
464+
`,
465+
`
466+
/**
467+
* Documentation for the function
468+
* @returns {string} A test value
469+
*/
470+
function example() {
471+
return 'test';
472+
}
473+
`
474+
)
475+
);
476+
});
477+
478+
test('should preserve inline and subsequent comments when removing import', async () => {
479+
const spec = new RecipeSpec();
480+
spec.recipe = fromVisitor(new RemoveImport("fs"));
481+
482+
//language=typescript
483+
await spec.rewriteRun(
484+
typescript(
485+
`
486+
import fs from 'fs'; // This import is unused
487+
488+
// Main application code starts here
489+
function main() {
490+
console.log('application');
491+
}
492+
`,
493+
`
494+
// Main application code starts here
495+
function main() {
496+
console.log('application');
497+
}
498+
`
499+
)
500+
);
501+
});
502+
503+
test('should preserve comments after multiple removed imports', async () => {
504+
const spec = new RecipeSpec();
505+
spec.recipe = fromVisitor(new RemoveImport("util", "isNullOrUndefined"));
506+
507+
//language=typescript
508+
await spec.rewriteRun(
509+
typescript(
510+
`
511+
import * as DiffGenerators from './diffgenerators/export';
512+
import { isNullOrUndefined } from 'util';
513+
514+
// Gets the xml files and passes them into diff generators
515+
class Foo {
516+
doSomething() {
517+
return DiffGenerators;
518+
}
519+
}
520+
`,
521+
`
522+
import * as DiffGenerators from './diffgenerators/export';
523+
524+
// Gets the xml files and passes them into diff generators
525+
class Foo {
526+
doSomething() {
527+
return DiffGenerators;
528+
}
529+
}
530+
`
531+
)
532+
);
533+
});
534+
});
416535
});

0 commit comments

Comments
 (0)
⚡