Codemod to wrap ignored rxjs subscriptions

rxjs-no-ignored-subscription.ts

import {
  API,
  ASTPath,
  BlockStatement,
  ClassBody,
  ClassDeclaration,
  ClassMethod,
  ClassProperty,
  Collection,
  ExpressionStatement,
  FileInfo,
  Identifier,
  ImportDeclaration,
  ImportSpecifier,
  JSCodeshift,
  Node,
} from 'jscodeshift';

const RXJS_LIBRARY_NAME = 'rxjs';
const RXJS_SUBSCRIPTION_TYPE = 'Subscription';
const DEFAULT_PROPERTY_NAME = 'subscription';
const UNKNOWN_DESTROY_METHOD_NAME = 'TODO_RENAME_AND_CALL_DESTROY_METHOD';
const UNKNOWN_DESTROY_METHOD_COMMENT = ` TODO: consider moving this to an existing method, or rename this and make sure it's called in the correct lifecycle method`;

export const parser = 'ts';

export default function transformer(file: FileInfo, api: API): string {
  if (!file.path.startsWith('client') || file.path.endsWith('.spec.ts')) {
    return null;
  }

  const j: JSCodeshift = api.jscodeshift;
  const root: Collection<Node> = j(file.source);
  const refactor: IgnoredSubscriptionRefactor = new IgnoredSubscriptionRefactor(j, root);

  if (refactor.ignoredSubscriptions.size() === 0) {
    return null;
  }

  if (!refactor.isSubscriptionImported) {
    refactor.addSubscriptionImport();
  }

  if (!refactor.isPropertyDeclared) {
    refactor.declareSubscriptionProperty();
    refactor.initializeSubscriptionProperty();
    refactor.unsubscribeSubscriptionProperty();
  }

  refactor.ignoredSubscriptions.forEach((path: ASTPath<ExpressionStatement>) => {
    refactor.wrapSubscriptionNode(path);
  });

  return root.toSource();
}

class IgnoredSubscriptionRefactor {
  constructor(
    private j: JSCodeshift,
    private root: Collection<Node>,
  ) { }

  private get subscriptionTypeName(): string | null {
    const j: JSCodeshift = this.j;
    const importPath: Collection<ImportSpecifier> = this.root
      .find(j.ImportDeclaration, {
        source: {
          type: 'StringLiteral',
          value: RXJS_LIBRARY_NAME,
        },
      })
      .find(j.ImportSpecifier, {
        imported: {
          type: 'Identifier',
          name: RXJS_SUBSCRIPTION_TYPE,
        },
      });

    if (importPath.length > 0) {
      const importNode: ImportSpecifier = importPath.get(0).node;
      return importNode.local?.name ?? importNode.imported.name;
    }

    return null;
  }

  private get propertyName(): string | null {
    const j: JSCodeshift = this.j;
    const propertyPath: Collection<Identifier> = this.root
      .find(j.ClassDeclaration)
      .find(j.ClassProperty, {
        typeAnnotation: {
          type: 'TSTypeAnnotation',
          typeAnnotation: {
            type: 'TSTypeReference',
            typeName: {
              name: this.targetSubscriptionTypeName,
            },
          },
        },
      })
      .find(j.Identifier);

    if (propertyPath.length > 0) {
      const propertyNode: Identifier = propertyPath.get(0).node;
      return propertyNode.name;
    }

    return null;
  }

  private get targetSubscriptionTypeName(): string {
    return this.subscriptionTypeName || RXJS_SUBSCRIPTION_TYPE;
  }

  private get targetPropertyName(): string {
    return this.propertyName || DEFAULT_PROPERTY_NAME;
  }

  private getExpressionMethodCalls(methodName: string): Collection<ExpressionStatement> {
    return this.root.find(this.j.ExpressionStatement, {
      expression: {
        type: 'CallExpression',
        callee: {
          type: 'MemberExpression',
          property: {
            type: 'Identifier',
            name: methodName,
          },
        },
      },
    });
  }

  private getLibraryImport(): Collection<ImportDeclaration> {
    return this.root.find(this.j.ImportDeclaration, {
      source: {
        type: 'StringLiteral',
        value: RXJS_LIBRARY_NAME,
      },
    });
  }

  private makeImportSpecifier(): ImportSpecifier {
    const j: JSCodeshift = this.j;
    return j.importSpecifier(j.identifier(RXJS_SUBSCRIPTION_TYPE));
  }

  private addImportDeclaration(): ImportDeclaration | null {
    const j: JSCodeshift = this.j;

    // find the latest external module import
    let latestModulePath: ASTPath<ImportDeclaration>;
    this.root
      .find(j.ImportDeclaration, {
        source: {
          type: 'StringLiteral',
        },
      })
      .forEach((path: ASTPath<ImportDeclaration>) => {
        const source: string = path.node.source.value.toString();
        if (!source.startsWith('.')) {
          latestModulePath = path;
        }
      });

    // no external export in this file so it may be tricky to figure where to place the import, so let this edge case be ignored
    // and the compilation will error due to the missing import which is easy to catch and fix later
    if (!latestModulePath) {
      return null;
    }

    // add new import declaration before the latest external module import
    const addedImportDeclaration: ImportDeclaration = j.importDeclaration(
      [this.makeImportSpecifier()], j.stringLiteral(RXJS_LIBRARY_NAME)
    );
    latestModulePath.insertBefore(
      j(addedImportDeclaration).toSource({ quote: 'single' }),
    );

    return addedImportDeclaration;
  }

  private addImportSpecifier(importDeclaration: ImportDeclaration): void {
    importDeclaration.specifiers.push(this.makeImportSpecifier());
  }

  private getTargetIndexForNgOnDestroy(path: ASTPath<ClassBody>): number {
    let ngOnInitIndex: number = path.node.body.findIndex((node: ClassMethod) => {
      return node.kind === 'method' && node.key.type === 'Identifier' && node.key.name === 'ngOnInit';
    });
    const lastIndex: number = path.node.body.length - 1;

    if (ngOnInitIndex !== -1) {
      ngOnInitIndex++;
    }

    const preferredIndexOrder: number[] = [ngOnInitIndex, lastIndex];
    return preferredIndexOrder.find((index: number) => index !== -1);
  }

  private isAngularComponent(): boolean {
    const j: JSCodeshift = this.j;
    return this.root
      .find(j.ClassDeclaration)
      .filter((path: ASTPath<any>) => {
        return path.node.decorators?.find((decorator: any) => {
          const decoratorName: string = decorator.expression.callee.name;
          return decoratorName === 'Component' || decoratorName === 'Directive';
        });
      })
      .size() > 0;
  }

  private getOrAddOnDestroyMethod(): ClassMethod {
    const j: JSCodeshift = this.j;
    const methodName: string = this.isAngularComponent() ? 'ngOnDestroy' : UNKNOWN_DESTROY_METHOD_NAME;
    const onDestroyMethod: Collection<ClassMethod> = this.root.find(j.ClassMethod, {
      key: {
        name: methodName,
      },
    });

    if (onDestroyMethod.size() > 0) {
      return onDestroyMethod.get(0).node;
    }

    const addedMethod: ClassMethod = j.classMethod(
      'method',
      j.identifier(methodName),
      [],
      j.blockStatement([])
    );
    addedMethod.access = 'public';
    addedMethod.returnType = j.tsTypeAnnotation(j.tsVoidKeyword());

    if (!this.isAngularComponent()) {
      addedMethod.comments = [j.commentLine(UNKNOWN_DESTROY_METHOD_COMMENT)];
    }

    this.root.find(j.ClassDeclaration).forEach((path: ASTPath<ClassDeclaration>) => {
      path.node.implements.push(
        j.tsExpressionWithTypeArguments(
          j.identifier('OnDestroy'),
        ) as any,
      );
    });

    this.root.find(j.ClassBody).forEach((path: ASTPath<ClassBody>) => {
      const classBodyIndex: number = this.getTargetIndexForNgOnDestroy(path);
      path.node.body.splice(classBodyIndex, 0, addedMethod);
    });

    return addedMethod;
  }

  get ignoredSubscriptions(): Collection<ExpressionStatement> {
    return this.getExpressionMethodCalls('subscribe');
  }

  get isSubscriptionImported(): boolean {
    return this.subscriptionTypeName !== null;
  }

  get isPropertyDeclared(): boolean {
    return this.propertyName !== null;
  }

  wrapSubscriptionNode(path: ASTPath<ExpressionStatement>): void {
    const j: JSCodeshift = this.j;
    const identifierName: string = `this.${this.targetPropertyName}.add`;

    j(path).replaceWith(
      j.expressionStatement(
        j.callExpression(j.identifier(identifierName), [path.node.expression]),
      )
    );
  }

  addSubscriptionImport(): void {
    const importDeclaration: Collection<ImportDeclaration> = this.getLibraryImport();

    if (importDeclaration.size() === 0) {
      this.addImportDeclaration();
    } else {
      const importNode: ImportDeclaration = importDeclaration.get(0).node;
      this.addImportSpecifier(importNode);
    }
  }

  declareSubscriptionProperty(): void {
    const j: JSCodeshift = this.j;
    const addedProperty: ClassProperty = j.classProperty(
      j.identifier(DEFAULT_PROPERTY_NAME),
      null,
      j.tsTypeAnnotation(j.tsTypeReference(j.identifier(this.targetSubscriptionTypeName))),
      false
    );
    addedProperty.access = 'private';

    const constructorMethod: Collection<ClassMethod> = this.root.find(j.ClassMethod, { kind: 'constructor' });
    const hasConstructor: boolean = constructorMethod.length !== 0;
    if (!hasConstructor) {
      addedProperty.value = j.newExpression(j.identifier(this.targetSubscriptionTypeName), []);
    }

    this.root.find(j.ClassBody).forEach((path: ASTPath<ClassBody>) => {
      path.node.body.unshift(addedProperty);
    });
  }

  initializeSubscriptionProperty(): void {
    const j: JSCodeshift = this.j;
    const constructorMethod: Collection<ClassMethod> = this.root.find(j.ClassMethod, {
      kind: 'constructor',
      body: {
        type: 'BlockStatement',
      },
    });

    if (constructorMethod.length === 0) {
      // edge case handled on declareSubscriptionProperty, the property initialization is inlined
      return;
    }

    const constructorBlock: BlockStatement = constructorMethod.find(j.BlockStatement).get(0).node;
    constructorBlock.body.unshift(
      j.expressionStatement(
        j.assignmentExpression(
          '=',
          j.memberExpression(
            j.thisExpression(),
            j.identifier(this.targetPropertyName),
          ),
          j.newExpression(j.identifier(this.targetSubscriptionTypeName), []),
        ),
      ),
    );
  }

  unsubscribeSubscriptionProperty(): void {
    const targetMethod: ClassMethod = this.getOrAddOnDestroyMethod();
    const identifierName: string = `this.${this.targetPropertyName}.unsubscribe`;

    const j: JSCodeshift = this.j;
    const unsubscribeStatement: ExpressionStatement = j.expressionStatement(
      j.callExpression(j.identifier(identifierName), []),
    );

    targetMethod.body.body.unshift(unsubscribeStatement);
  }
}

TODO: try to incorporate the helper couldBeObservable used at https://github.com/cartant/eslint-plugin-rxjs/blob/main/source/rules/no-ignored-subscription.ts