fix(migrations): cf migration - improve import declaration handling (#53622)

thePunderWoman · thePunderWoman · commit 7ac60bab9a4b · 2023-12-18T22:10:58.000Z
This should make the import declaration symbol removal a bit more robust and handle more than just CommonModule safely. PR Close #53622
diff --git a/packages/core/schematics/ng-generate/control-flow-migration/identifier-lookup.ts b/packages/core/schematics/ng-generate/control-flow-migration/identifier-lookup.ts
@@ -9,10 +9,10 @@
 import ts from 'typescript';
 
 export function lookupIdentifiersInSourceFile(
-    sourceFile: ts.SourceFile, name: string): Set<ts.Identifier> {
+    sourceFile: ts.SourceFile, names: string[]): Set<ts.Identifier> {
   const results = new Set<ts.Identifier>();
   const visit = (node: ts.Node): void => {
-    if (ts.isIdentifier(node) && node.text === name) {
+    if (ts.isIdentifier(node) && names.includes(node.text)) {
       results.add(node);
     }
     ts.forEachChild(node, visit);
diff --git a/packages/core/schematics/ng-generate/control-flow-migration/migration.ts b/packages/core/schematics/ng-generate/control-flow-migration/migration.ts
@@ -82,6 +82,7 @@ export function migrateTemplate(
       // and in that case we can't safely remove the common module import.
       componentFile.verifyCanRemoveImports();
     }
+    file.verifyCanRemoveImports();
 
     errors = [
       ...ifResult.errors,
diff --git a/packages/core/schematics/ng-generate/control-flow-migration/types.ts b/packages/core/schematics/ng-generate/control-flow-migration/types.ts
@@ -23,6 +23,13 @@ export const endMarker = '✢';
 export const startI18nMarker = '⚈';
 export const endI18nMarker = '⚉';
 
+export const importRemovals = [
+  'NgIf', 'NgIfElse', 'NgIfThenElse', 'NgFor', 'NgForOf', 'NgForTrackBy', 'NgSwitch',
+  'NgSwitchCase', 'NgSwitchDefault'
+];
+
+export const importWithCommonRemovals = [...importRemovals, 'CommonModule'];
+
 function allFormsOf(selector: string): string[] {
   return [
     selector,
@@ -308,7 +315,7 @@ export class AnalyzedFile {
     // skip this check entirely
     if (this.removeCommonModule) {
       const importDeclaration = this.importRanges.find(r => r.type === 'importDeclaration');
-      const instances = lookupIdentifiersInSourceFile(this.sourceFile, 'CommonModule');
+      const instances = lookupIdentifiersInSourceFile(this.sourceFile, importWithCommonRemovals);
       let foundImportDeclaration = false;
       let count = 0;
       for (let range of this.importRanges) {
diff --git a/packages/core/schematics/ng-generate/control-flow-migration/util.ts b/packages/core/schematics/ng-generate/control-flow-migration/util.ts
@@ -10,13 +10,8 @@ import {Attribute, Element, HtmlParser, Node, ParseTreeResult, visitAll} from '@
 import {dirname, join} from 'path';
 import ts from 'typescript';
 
-import {AnalyzedFile, CommonCollector, ElementCollector, ElementToMigrate, endI18nMarker, endMarker, i18nCollector, ParseResult, startI18nMarker, startMarker, Template, TemplateCollector} from './types';
+import {AnalyzedFile, CommonCollector, ElementCollector, ElementToMigrate, endI18nMarker, endMarker, i18nCollector, importRemovals, importWithCommonRemovals, ParseResult, startI18nMarker, startMarker, Template, TemplateCollector} from './types';
 
-const importRemovals = [
-  'NgIf', 'NgIfElse', 'NgIfThenElse', 'NgFor', 'NgForOf', 'NgForTrackBy', 'NgSwitch',
-  'NgSwitchCase', 'NgSwitchDefault'
-];
-const importWithCommonRemovals = [...importRemovals, 'CommonModule'];
 const startMarkerRegex = new RegExp(startMarker, 'gm');
 const endMarkerRegex = new RegExp(endMarker, 'gm');
 const startI18nMarkerRegex = new RegExp(startI18nMarker, 'gm');
diff --git a/packages/core/schematics/test/control_flow_migration_spec.ts b/packages/core/schematics/test/control_flow_migration_spec.ts
@@ -5087,6 +5087,68 @@ describe('control flow migration', () => {
 
       expect(actual).toBe(expected);
     });
+
+    it('should not remove imports when mismatch in counts', async () => {
+      writeFile('/comp.ts', [
+        `import {CommonModule} from '@angular/common';`,
+        `import {Component, NgModule, Pipe, PipeTransform} from '@angular/core';`,
+        `@Component({`,
+        `  selector: 'description',`,
+        `  template: \`<span>{{getDescription()}}</span>\`,`,
+        `})`,
+        `export class DescriptionController {`,
+        `  getDescription(): string {`,
+        `    return 'stuff';`,
+        `  }`,
+        `}`,
+        ``,
+        `@Pipe({name: 'description'})`,
+        `export class DescriptionPipe implements PipeTransform {`,
+        `  transform(nameString?: string): string {`,
+        `    return nameString ?? '';`,
+        `  }`,
+        `}`,
+        `@NgModule({`,
+        `  declarations: [DescriptionController, DescriptionPipe],`,
+        `  imports: [CommonModule],`,
+        `  providers: [],`,
+        `  exports: [DescriptionController, DescriptionPipe],`,
+        `})`,
+        `export class DescriptionModule {}`,
+      ].join('\n'));
+
+      await runMigration();
+      const actual = tree.readContent('/comp.ts');
+      const expected = [
+        `import {CommonModule} from '@angular/common';`,
+        `import {Component, NgModule, Pipe, PipeTransform} from '@angular/core';`,
+        `@Component({`,
+        `  selector: 'description',`,
+        `  template: \`<span>{{getDescription()}}</span>\`,`,
+        `})`,
+        `export class DescriptionController {`,
+        `  getDescription(): string {`,
+        `    return 'stuff';`,
+        `  }`,
+        `}`,
+        ``,
+        `@Pipe({name: 'description'})`,
+        `export class DescriptionPipe implements PipeTransform {`,
+        `  transform(nameString?: string): string {`,
+        `    return nameString ?? '';`,
+        `  }`,
+        `}`,
+        `@NgModule({`,
+        `  declarations: [DescriptionController, DescriptionPipe],`,
+        `  imports: [CommonModule],`,
+        `  providers: [],`,
+        `  exports: [DescriptionController, DescriptionPipe],`,
+        `})`,
+        `export class DescriptionModule {}`,
+      ].join('\n');
+
+      expect(actual).toBe(expected);
+    });
   });
 
   describe('no migration needed', () => {

Original file line number	Diff line number	Diff line change
`@@ -82,6 +82,7 @@ export function migrateTemplate(`
`82`	`82`	`// and in that case we can't safely remove the common module import.`
`83`	`83`	`componentFile.verifyCanRemoveImports();`
`84`	`84`	`}`
	`85`	`+ file.verifyCanRemoveImports();`
`85`	`86`
`86`	`87`	`errors = [`
`87`	`88`	`...ifResult.errors,`