A React codemod that will transform proptypes to typescript type alias for component types

proptypes-to-typescript.ts

// Modified version of https://github.com/mskelton/ratchet

import type { NodePath } from 'ast-types/lib/node-path';
import type {
  API,
  Collection,
  CommentBlock,
  CommentLine,
  FileInfo,
  Identifier,
  JSCodeshift,
  Literal,
  TSAnyKeyword,
  TSFunctionType,
} from 'jscodeshift';

let j: JSCodeshift;

function reactType(type: string) {
  return j.tsQualifiedName(j.identifier('React'), j.identifier(type));
}

type TSType = {
  comments: (CommentLine | CommentBlock)[];
  key: Identifier | Literal;
  required: boolean;
  type: TSAnyKeyword | TSFunctionType;
};

function createPropertySignature({ comments, key, required, type }: TSType) {
  return j.tsPropertySignature.from({
    comments,
    key,
    optional: !required,
    typeAnnotation: j.tsTypeAnnotation(type),
  });
}

function isCustomValidator(path: NodePath) {
  return path.get('type').value === 'FunctionExpression' || path.get('type').value === 'ArrowFunctionExpression';
}

const resolveRequired = (path: NodePath) => (isRequired(path) ? path.get('object') : path);

function getTSType(path: NodePath) {
  const { value: name } =
    path.get('type').value === 'MemberExpression'
      ? path.get('property', 'name')
      : path.get('callee', 'property', 'name');

  switch (name) {
    case 'func': {
      const restElement = j.restElement.from({
        argument: j.identifier('args'),
        typeAnnotation: j.tsTypeAnnotation(j.tsArrayType(j.tsAnyKeyword())),
      });

      return j.tsFunctionType.from({
        parameters: [restElement],
        typeAnnotation: j.tsTypeAnnotation(j.tsVoidKeyword()),
      });
    }

    case 'arrayOf': {
      const type = path.get('arguments', 0);
      return isCustomValidator(type) ? j.tsAnyKeyword() : j.tsArrayType(getTSType(resolveRequired(type)));
    }

    case 'objectOf': {
      const type = path.get('arguments', 0);
      return isCustomValidator(type)
        ? j.tsAnyKeyword()
        : j.tsTypeReference(
          j.identifier('Record'),
          j.tsTypeParameterInstantiation([j.tsStringKeyword(), getTSType(resolveRequired(type))])
        );
    }

    case 'oneOf': {
      const arg = path.get('arguments', 0);

      return arg.get('type').value !== 'ArrayExpression'
        ? j.tsArrayType(j.tsAnyKeyword())
        : j.tsUnionType(arg.get('elements').value.map(({ value }) => j.tsLiteralType(j.stringLiteral(value))));
    }

    case 'oneOfType':
      return j.tsUnionType(path.get('arguments', 0, 'elements').map(getTSType));

    case 'instanceOf':
      return j.tsTypeReference(j.identifier(path.get('arguments', 0, 'name').value));

    case 'shape':
    case 'exact':
      return j.tsTypeLiteral(path.get('arguments', 0, 'properties').map(mapType).map(createPropertySignature));
  }

  const map = {
    any: j.tsAnyKeyword(),
    array: j.tsArrayType(j.tsAnyKeyword()),
    bool: j.tsBooleanKeyword(),
    element: j.tsTypeReference(reactType('ReactElement')),
    elementType: j.tsTypeReference(reactType('ElementType')),
    node: j.tsTypeReference(reactType('ReactNode')),
    number: j.tsNumberKeyword(),
    object: j.tsAnyKeyword(),
    string: j.tsStringKeyword(),
    symbol: j.tsSymbolKeyword(),
  };

  return map[name] || j.tsAnyKeyword();
}

const isRequired = (path: NodePath) =>
  path.get('type').value === 'MemberExpression' && path.get('property', 'name').value === 'isRequired';

function mapType(path: NodePath): TSType {
  const required = isRequired(path.get('value'));
  const key = path.get('key').value;
  const comments = path.get('leadingComments').value;
  const type = getTSType(required ? path.get('value', 'object') : path.get('value'));

  path.replace();

  return {
    comments: comments ?? [],
    key,
    required,
    type,
  };
}

type CollectedTypes = {
  component: string;
  types: TSType[];
}[];

function getTSTypes(source: Collection, _getComponentName: (path: NodePath) => string) {
  const collected = [] as CollectedTypes;

  source
    .filter(path => path.value)
    .forEach(path => {
      collected.push({
        component: _getComponentName(path),
        types: path
          .filter(({ value }) => value.type === 'ObjectProperty' || value.type === 'ObjectMethod', null)
          .map(mapType, null),
      });
    });

  return collected;
}

function getFunctionParent(path: NodePath) {
  return path.parent.get('type').value === 'Program' ? path : getFunctionParent(path.parent);
}

function getComponentName(path: NodePath) {
  const root = path.get('type').value === 'ArrowFunctionExpression' ? path.parent : path;

  return root.get('id', 'name').value;
}

function createTypeAlias(path: NodePath, componentTypes: CollectedTypes) {
  const componentName = getComponentName(path);
  const types = componentTypes.find(t => t.component === componentName);

  // If the component doesn't have propTypes, ignore it
  if (!types) return;

  const typeName = `${componentName}Props`;

  // Add the TS types before the function/class
  getFunctionParent(path).insertBefore(
    j.tsTypeAliasDeclaration(j.identifier(typeName), j.tsTypeLiteral(types.types.map(createPropertySignature)))
  );

  return typeName;
}

function addFunctionTSTypes(source: Collection, componentTypes: CollectedTypes) {
  source.forEach(path => {
    const typeName = createTypeAlias(path, componentTypes);
    if (!typeName) return;

    // Add the TS types to the props param
    path.get('params', 0).value.typeAnnotation = j.tsTypeReference(
      // For some reason, jscodeshift isn't adding the colon so we have to do
      // that ourselves.
      j.identifier(`: ${typeName}`)
    );
  });
}

function addClassTSTypes(source: Collection, componentTypes: CollectedTypes) {
  source.find(j.ClassDeclaration).forEach(path => {
    const typeName = createTypeAlias(path, componentTypes);
    if (!typeName) return;

    // Add the TS types to the React.Component super class
    path.value.superTypeParameters = j.tsTypeParameterInstantiation([j.tsTypeReference(j.identifier(typeName))]);
  });
}

function collectPropTypes(source: Collection) {
  return source
    .find(j.AssignmentExpression)
    .filter(path => path.get('left', 'property', 'name').value === 'propTypes')
    .map(path => path.get('right', 'properties'));
}

function collectStaticPropTypes(source: Collection) {
  return source
    .find(j.ClassProperty)
    .filter(path => !!path.value.static)
    .filter(path => path.get('key', 'name').value === 'propTypes')
    .map(path => path.get('value', 'properties'));
}

function cleanup(source: Collection, propTypes: Collection, staticPropTypes: Collection) {
  propTypes.forEach(path => {
    if (!path.parent.get('right', 'properties', 'length').value) {
      path.parent.prune();
    }
  });

  staticPropTypes.forEach(path => {
    if (!path.parent.get('value', 'properties', 'length').value) {
      path.parent.prune();
    }
  });

  const propTypesUsages = source
    .find(j.MemberExpression)
    .filter(path => path.get('object', 'name').value === 'PropTypes');

  // We can remove the import without caring about the preserve-prop-types
  // option since the criteria for removal is that no PropTypes.* member
  // expressions exist.
  if (propTypesUsages.length === 0) {
    source
      .find(j.ImportDeclaration)
      .filter(path => path.value.source.value === 'prop-types')
      .remove();
  }

  propTypes.remove();
  staticPropTypes.remove();
}

// Use the TSX to allow parsing of TypeScript code that still contains prop
// types. Though not typical, this exists in the wild.
export const parser = 'tsx';

export default function (file: FileInfo, api: API) {
  j = api.jscodeshift;
  const source = j(file.source);

  const propTypes = collectPropTypes(source);
  const tsTypes = getTSTypes(propTypes, path => path.parent.get('left', 'object', 'name').value);

  const staticPropTypes = collectStaticPropTypes(source);
  const staticTSTypes = getTSTypes(staticPropTypes, path => path.parent.parent.parent.value.id.name);

  addFunctionTSTypes(source.find(j.FunctionDeclaration), tsTypes);
  addFunctionTSTypes(source.find(j.FunctionExpression), tsTypes);
  addFunctionTSTypes(source.find(j.ArrowFunctionExpression), tsTypes);
  addClassTSTypes(source, tsTypes);
  addClassTSTypes(source, staticTSTypes);

  // Remove empty propTypes expressions and imports
  cleanup(source, propTypes, staticPropTypes);

  return source.toSource();
}