[core] Handle refType in the proptypes generation

mui · Sep 11, 2023 · f9b8721 · f9b8721
1 parent cda1ffe
commit f9b8721
Show file tree

Hide file tree

Showing 6 changed files with 114 additions and 39 deletions.
diff --git a/packages/typescript-to-proptypes/src/createType.ts b/packages/typescript-to-proptypes/src/createType.ts
@@ -15,6 +15,7 @@ import {
   StringType,
   ObjectType,
   NumericType,
+  ReactRefType,
 } from './models';
 
 export function createAnyType(init: { jsDoc: string | undefined }): AnyType {
@@ -64,6 +65,13 @@ export function createElementType(init: {
   };
 }
 
+export function createReactRefType(init: { jsDoc: string | undefined }): ReactRefType {
+  return {
+    type: 'ReactRefNode',
+    jsDoc: init.jsDoc,
+  };
+}
+
 export function createFunctionType(init: { jsDoc: string | undefined }): FunctionType {
   return {
     type: 'FunctionNode',

diff --git a/packages/typescript-to-proptypes/src/generatePropTypes.ts b/packages/typescript-to-proptypes/src/generatePropTypes.ts
@@ -155,6 +155,10 @@ export function generatePropTypes(
       return `${importedName}.bool`;
     }
 
+    if (propType.type === 'ReactRefNode') {
+      return 'refType';
+    }
+
     if (propType.type === 'NumericNode') {
       return `${importedName}.number`;
     }

diff --git a/packages/typescript-to-proptypes/src/getPropTypesFromFile.ts b/packages/typescript-to-proptypes/src/getPropTypesFromFile.ts
@@ -19,6 +19,7 @@ import {
   createNumericType,
   createObjectType,
   createStringType,
+  createReactRefType,
 } from './createType';
 import { PropTypeDefinition, PropTypesComponent, PropType } from './models';
 
@@ -291,6 +292,23 @@ function checkSymbol({
     ts.isTypeReferenceNode(declaration.type)
   ) {
     const name = declaration.type.typeName.getText();
+
+    if (name === 'React.Ref') {
+      return {
+        $$id: project.createPropTypeId(symbol),
+        name: symbol.getName(),
+        jsDoc,
+        filenames: symbolFilenames,
+        propType: createUnionType({
+          jsDoc,
+          types: [
+            createUndefinedType({ jsDoc: undefined }),
+            createReactRefType({ jsDoc: undefined }),
+          ],
+        }),
+      };
+    }
+
     if (
       name === 'React.ElementType' ||
       name === 'React.JSXElementConstructor' ||

diff --git a/packages/typescript-to-proptypes/src/injectPropTypesInFile.ts b/packages/typescript-to-proptypes/src/injectPropTypesInFile.ts
@@ -2,7 +2,7 @@ import * as babel from '@babel/core';
 import * as babelTypes from '@babel/types';
 import { v4 as uuid } from 'uuid';
 import { generatePropTypes, GeneratePropTypesOptions } from './generatePropTypes';
-import { PropTypesComponent, PropTypeDefinition, LiteralType } from './models';
+import { PropTypesComponent, PropTypeDefinition, LiteralType, PropType } from './models';
 
 export interface InjectPropTypesInFileOptions
   extends Pick<
@@ -170,9 +170,11 @@ function createBabelPlugin({
     return includeUnusedProps ? true : data.usedProps.includes(data.prop.name);
   };
 
-  let importName = '';
-  let needImport = false;
-  let alreadyImported = false;
+  let importNodeFromMuiUtils: babel.NodePath | null = null;
+  let hasRefTypeImportFromMuiUtils = false;
+  let needRefTypeImportFromMuiUtils = false;
+  let alreadyImportedPropTypesName: string | null = null;
+  let needImportFromPropTypePackage = false;
   let originalPropTypesPath: null | babel.NodePath = null;
   const previousPropTypesSource = new Map<string, string>();
 
@@ -186,15 +188,21 @@ function createBabelPlugin({
 
     const source = generatePropTypes(props, {
       ...otherOptions,
-      importedName: importName,
+      importedName: alreadyImportedPropTypesName ?? 'PropTypes',
       previousPropTypesSource,
       reconcilePropTypes,
       shouldInclude: (prop) => shouldInclude({ component: props, prop, usedProps }),
     });
     const emptyPropTypes = source === '';
 
     if (!emptyPropTypes) {
-      needImport = true;
+      needImportFromPropTypePackage = true;
+    }
+
+    // TODO: If we filter the injected proptypes in this file instead of in generatePropTypes.
+    // Then we could check correctly if some prop-type is a ref and not just rely on some string analysis.
+    if (source.includes(': refType')) {
+      needRefTypeImportFromMuiUtils = true;
     }
 
     const placeholder = `const a${uuid().replace(/-/g, '_')} = null;`;
@@ -233,25 +241,23 @@ function createBabelPlugin({
     visitor: {
       Program: {
         enter(path, state: any) {
-          if (
-            !path.node.body.some((n) => {
-              if (
-                babelTypes.isImportDeclaration(n) &&
-                n.source.value === 'prop-types' &&
-                n.specifiers.length
-              ) {
-                importName = n.specifiers[0].local.name;
-                alreadyImported = true;
-                return true;
-              }
-              return false;
-            })
-          ) {
-            importName = 'PropTypes';
-          }
-
           path.get('body').forEach((nodePath) => {
             const { node } = nodePath;
+
+            if (babelTypes.isImportDeclaration(node) && node.specifiers.length) {
+              if (node.source.value === 'prop-types') {
+                alreadyImportedPropTypesName = node.specifiers[0].local.name;
+              } else if (node.source.value === '@mui/utils') {
+                importNodeFromMuiUtils = nodePath;
+                const specifier = node.specifiers.find(
+                  (el) => babelTypes.isImportSpecifier(el) && el.local.name === 'refType',
+                );
+                if (specifier) {
+                  hasRefTypeImportFromMuiUtils = true;
+                }
+              }
+            }
+
             if (
               babelTypes.isExpressionStatement(node) &&
               babelTypes.isAssignmentExpression(node.expression, { operator: '=' }) &&
@@ -293,23 +299,54 @@ function createBabelPlugin({
           });
         },
         exit(path) {
-          if (alreadyImported || !needImport) {
-            return;
+          if (alreadyImportedPropTypesName == null && needImportFromPropTypePackage) {
+            const propTypesImport = babel.template.ast(
+              `import PropTypes from 'prop-types'`,
+            ) as babel.types.ImportDeclaration;
+
+            const firstImport = path
+              .get('body')
+              .find((nodePath) => babelTypes.isImportDeclaration(nodePath.node));
+
+            // Insert import after the first one to avoid issues with comment flags
+            if (firstImport) {
+              firstImport.insertAfter(propTypesImport);
+            } else {
+              path.node.body = [propTypesImport, ...path.node.body];
+            }
           }
 
-          const propTypesImport = babel.template.ast(
-            `import ${importName} from 'prop-types'`,
-          ) as babel.types.ImportDeclaration;
-
-          const firstImport = path
-            .get('body')
-            .find((nodePath) => babelTypes.isImportDeclaration(nodePath.node));
-
-          // Insert import after the first one to avoid issues with comment flags
-          if (firstImport) {
-            firstImport.insertAfter(propTypesImport);
-          } else {
-            path.node.body = [propTypesImport, ...path.node.body];
+          if (needRefTypeImportFromMuiUtils && !hasRefTypeImportFromMuiUtils) {
+            if (importNodeFromMuiUtils) {
+              const node = importNodeFromMuiUtils.node as babel.types.ImportDeclaration;
+              importNodeFromMuiUtils.replaceWith(
+                babelTypes.importDeclaration(
+                  [
+                    ...node.specifiers,
+                    babelTypes.importSpecifier(
+                      babelTypes.identifier('refType'),
+                      babelTypes.identifier('refType'),
+                    ),
+                  ],
+                  babelTypes.stringLiteral('@mui/utils'),
+                ),
+              );
+            } else {
+              const refTypeImport = babel.template.ast(
+                `import { refType } from '@mui/utils'`,
+              ) as babel.types.ImportDeclaration;
+
+              const firstImport = path
+                .get('body')
+                .find((nodePath) => babelTypes.isImportDeclaration(nodePath.node));
+
+              // Insert import after the first one to avoid issues with comment flags
+              if (firstImport) {
+                firstImport.insertAfter(refTypeImport);
+              } else {
+                path.node.body = [refTypeImport, ...path.node.body];
+              }
+            }
           }
         },
       },

diff --git a/packages/typescript-to-proptypes/src/models.ts b/packages/typescript-to-proptypes/src/models.ts
@@ -21,6 +21,7 @@ export type PropType =
   | BooleanType
   | DOMElementType
   | ElementType
+  | ReactRefType
   | FunctionType
   | InstanceOfType
   | InterfaceType
@@ -63,6 +64,10 @@ export interface ElementType extends BasePropType {
   type: 'ElementNode';
 }
 
+export interface ReactRefType extends BasePropType {
+  type: 'ReactRefNode';
+}
+
 export interface FunctionType extends BasePropType {
   type: 'FunctionNode';
 }

diff --git a/scripts/generateProptypes.ts b/scripts/generateProptypes.ts
@@ -248,7 +248,10 @@ async function generateProptypes(
       ensureBabelPluginTransformReactRemovePropTypesIntegration: true,
       getSortLiteralUnions,
       reconcilePropTypes: (prop, previous, generated) => {
-        const usedCustomValidator = previous !== undefined && !previous.startsWith('PropTypes');
+        const usedCustomValidator =
+          previous !== undefined &&
+          !previous.startsWith('PropTypes') &&
+          !previous.startsWith('refType');
         const ignoreGenerated =
           previous !== undefined &&
           previous.startsWith('PropTypes /* @typescript-to-proptypes-ignore */');