fix: replace existing propTypes when removeExistingPropTypes (#15)

src/injector.ts

@@ -100,6 +100,7 @@ function plugin(
 	let importName = '';
 	let needImport = false;
 	let alreadyImported = false;
+	let originalPropTypesPath: null | babel.NodePath = null;
 
 	return {
 		visitor: {
@@ -121,19 +122,17 @@ function plugin(
 						importName = 'PropTypes';
 					}
 
-					if (removeExistingPropTypes) {
-						path.get('body').forEach((nodePath) => {
-							const { node } = nodePath;
-							if (
-								babelTypes.isExpressionStatement(node) &&
-								babelTypes.isAssignmentExpression(node.expression, { operator: '=' }) &&
-								babelTypes.isMemberExpression(node.expression.left) &&
-								babelTypes.isIdentifier(node.expression.left.property, { name: 'propTypes' })
-							) {
-								nodePath.remove();
-							}
-						});
-					}
+					path.get('body').forEach((nodePath) => {
+						const { node } = nodePath;
+						if (
+							babelTypes.isExpressionStatement(node) &&
+							babelTypes.isAssignmentExpression(node.expression, { operator: '=' }) &&
+							babelTypes.isMemberExpression(node.expression.left) &&
+							babelTypes.isIdentifier(node.expression.left.property, { name: 'propTypes' })
+						) {
+							originalPropTypesPath = nodePath;
+						}
+					});
 				},
 				exit(path) {
 					if (alreadyImported || !needImport) return;
@@ -280,8 +279,9 @@ function plugin(
 
 		mapOfPropTypes.set(placeholder, source);
 
-		// Insert prop types
-		if (babelTypes.isExportNamedDeclaration(path.parent)) {
+		if (removeExistingPropTypes && originalPropTypesPath !== null) {
+			originalPropTypesPath.replaceWith(babel.template.ast(placeholder) as babelTypes.Statement);
+		} else if (babelTypes.isExportNamedDeclaration(path.parent)) {
 			path.insertAfter(babel.template.ast(`export { ${nodeName} };`));
 			path.insertAfter(babel.template.ast(placeholder));
 			path.parentPath.replaceWith(path.node);

test/code-order/input.d.ts

@@ -0,0 +1,7 @@
+import * as React from 'react';
+
+export interface Props {
+	value: unknown;
+}
+
+export default function Component(props: Props): JSX.Element;

test/code-order/input.js

@@ -0,0 +1,15 @@
+import React from 'react';
+import PropTypes from 'prop-types';
+
+function Component(props) {
+	const { value } = props;
+	return <div>{value}</div>;
+}
+
+const someValidator = () => new Error();
+
+Component.propTypes = {
+	value: PropTypes.any,
+};
+
+export default Component;

test/code-order/options.ts

@@ -0,0 +1,9 @@
+import { TestOptions } from '../types';
+
+const options: TestOptions = {
+	injector: {
+		removeExistingPropTypes: true,
+	},
+};
+
+export default options;

test/code-order/output.js

@@ -0,0 +1,15 @@
+import React from 'react';
+import PropTypes from 'prop-types';
+
+function Component(props) {
+	const { value } = props;
+	return <div>{value}</div>;
+}
+
+const someValidator = () => new Error();
+
+Component.propTypes = {
+	value: PropTypes.any.isRequired,
+};
+
+export default Component;

test/code-order/output.json

@@ -0,0 +1,10 @@
+{
+	"type": "ProgramNode",
+	"body": [
+		{
+			"type": "ComponentNode",
+			"name": "Component",
+			"types": [{ "type": "PropTypeNode", "name": "value", "propType": { "type": "AnyNode" } }]
+		}
+	]
+}

test/index.test.ts

@@ -22,6 +22,7 @@ for (const testCase of testCases) {
 	const astPath = path.join(dirname, 'output.json');
 	const outputPath = path.join(dirname, 'output.js');
 	const optionsPath = path.join(dirname, 'options.ts');
+	const inputJS = path.join(dirname, 'input.js');
 
 	it(testName, () => {
 		const options: TestOptions = fs.existsSync(optionsPath) ? require(optionsPath).default : {};
@@ -42,21 +43,28 @@ for (const testCase of testCases) {
 		}
 		//#endregion
 
-		let result = '';
-		// For d.ts files we just generate the AST
+		let inputSource = null;
 		if (testCase.endsWith('.d.ts')) {
-			result = ttp.generate(ast, options.generator);
-		}
-		// For .tsx? files we transpile them and inject the proptypes
-		else {
-			const transpiled = ttp.ts.transpileModule(fs.readFileSync(testCase, 'utf8'), {
+			try {
+				inputSource = fs.readFileSync(inputJS, { encoding: 'utf8' });
+			} catch (error) {}
+		} else {
+			inputSource = ttp.ts.transpileModule(fs.readFileSync(testCase, 'utf8'), {
 				compilerOptions: {
 					target: ttp.ts.ScriptTarget.ESNext,
 					jsx: ttp.ts.JsxEmit.Preserve,
 				},
 			}).outputText;
+		}
 
-			const injected = ttp.inject(ast, transpiled, options.injector);
+		let result = '';
+		// For d.ts files we just generate the AST
+		if (!inputSource) {
+			result = ttp.generate(ast, options.generator);
+		}
+		// For .tsx? files we transpile them and inject the proptypes
+		else {
+			const injected = ttp.inject(ast, inputSource, options.injector);
 			if (!injected) {
 				throw new Error('Injection failed');
 			}