[react-transform] Support more component definitions & expand tests

.changeset/sour-bears-complain.md

@@ -0,0 +1,5 @@
+---
+"@preact/signals-react-transform": patch
+---
+
+Add support for auto-transforming more ways to specify components: object methods, member assignments, export default components, components wrapped in HoCs like memo and forwardRef

packages/react-transform/package.json

@@ -59,12 +59,14 @@
     "@types/babel__core": "^7.20.1",
     "@types/babel__helper-module-imports": "^7.18.0",
     "@types/babel__helper-plugin-utils": "^7.10.0",
+    "@types/prettier": "^2.7.3",
     "@types/react": "^18.0.18",
     "@types/react-dom": "^18.0.6",
     "@types/use-sync-external-store": "^0.0.3",
     "assert": "^2.0.0",
     "buffer": "^6.0.3",
     "path": "^0.12.7",
+    "prettier": "^2.7.1",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
     "react-router-dom": "^6.9.0"

packages/react-transform/src/index.ts

@@ -8,10 +8,6 @@ import {
 } from "@babel/core";
 import { isModule, addNamed } from "@babel/helper-module-imports";
 
-// TODO:
-// - how to trigger rerenders on attributes change if transform never sees
-//   `.value`?
-
 interface PluginArgs {
 	types: typeof BabelTypes;
 	template: typeof BabelTemplate;
@@ -51,35 +47,125 @@ function setOnFunctionScope(path: NodePath, key: string, value: any) {
 type FunctionLike =
 	| BabelTypes.ArrowFunctionExpression
 	| BabelTypes.FunctionExpression
-	| BabelTypes.FunctionDeclaration;
-
-function testFunctionName<T extends FunctionLike>(
-	predicate: (name: string | null) => boolean
-): (path: NodePath<T>) => boolean {
-	return (path: NodePath<T>) => {
-		if (
-			path.node.type === "ArrowFunctionExpression" ||
-			path.node.type === "FunctionExpression"
-		) {
-			return (
-				path.parentPath.node.type === "VariableDeclarator" &&
-				path.parentPath.node.id.type === "Identifier" &&
-				predicate(path.parentPath.node.id.name)
-			);
-		} else if (path.node.type === "FunctionDeclaration") {
-			return predicate(path.node.id?.name ?? null);
+	| BabelTypes.FunctionDeclaration
+	| BabelTypes.ObjectMethod;
+
+/**
+ * Simple "best effort" to get the base name of a file path. Not fool proof but
+ * works in browsers and servers. Good enough for our purposes.
+ */
+function basename(filename: string | undefined): string | undefined {
+	return filename?.split(/[\\/]/).pop();
+}
+
+const DefaultExportSymbol = Symbol("DefaultExportSymbol");
+
+/**
+ * If the function node has a name (i.e. is a function declaration with a
+ * name), return that. Else return null.
+ */
+function getFunctionNodeName(path: NodePath<FunctionLike>): string | null {
+	if (path.node.type === "FunctionDeclaration" && path.node.id) {
+		return path.node.id.name;
+	} else if (path.node.type === "ObjectMethod") {
+		if (path.node.key.type === "Identifier") {
+			return path.node.key.name;
+		} else if (path.node.key.type === "StringLiteral") {
+			return path.node.key.value;
+		}
+	}
+
+	return null;
+}
+
+/**
+ * Given a function path's parent path, determine the "name" associated with the
+ * function. If the function is an inline default export (e.g. `export default
+ * () => {}`), returns a symbol indicating it is a default export. If the
+ * function is an anonymous function wrapped in higher order functions (e.g.
+ * memo(() => {})) we'll climb through the higher order functions to find the
+ * name of the variable that the function is assigned to, if any. Other cases
+ * handled too (see implementation). Else returns null.
+ */
+function getFunctionNameFromParent(
+	parentPath: NodePath<BabelTypes.Node>
+): string | null | typeof DefaultExportSymbol {
+	if (
+		parentPath.node.type === "VariableDeclarator" &&
+		parentPath.node.id.type === "Identifier"
+	) {
+		return parentPath.node.id.name;
+	} else if (parentPath.node.type === "AssignmentExpression") {
+		const left = parentPath.node.left;
+		if (left.type === "Identifier") {
+			return left.name;
+		} else if (left.type === "MemberExpression") {
+			let property = left.property;
+			while (property.type === "MemberExpression") {
+				property = property.property;
+			}
+
+			if (property.type === "Identifier") {
+				return property.name;
+			} else if (property.type === "StringLiteral") {
+				return property.value;
+			}
+
+			return null;
 		} else {
-			return false;
+			return null;
 		}
-	};
+	} else if (parentPath.node.type === "ExportDefaultDeclaration") {
+		return DefaultExportSymbol;
+	} else if (
+		parentPath.node.type === "CallExpression" &&
+		parentPath.parentPath != null
+	) {
+		// If our parent is a Call Expression, then this function expression is
+		// wrapped in some higher order functions. Recurse through the higher order
+		// functions to determine if this expression is assigned to a name we can
+		// use as the function name
+		return getFunctionNameFromParent(parentPath.parentPath);
+	} else {
+		return null;
+	}
+}
+
+/* Determine the name of a function */
+function getFunctionName(
+	path: NodePath<FunctionLike>
+): string | typeof DefaultExportSymbol | null {
+	let nodeName = getFunctionNodeName(path);
+	if (nodeName) {
+		return nodeName;
+	}
+
+	return getFunctionNameFromParent(path.parentPath);
+}
+
+function fnNameStartsWithCapital(
+	path: NodePath<FunctionLike>,
+	filename: string | undefined
+): boolean {
+	const name = getFunctionName(path);
+	if (!name) return false;
+	if (name === DefaultExportSymbol) {
+		return basename(filename)?.match(/^[A-Z]/) != null ?? false;
+	}
+	return name.match(/^[A-Z]/) != null;
 }
+function fnNameStartsWithUse(
+	path: NodePath<FunctionLike>,
+	filename: string | undefined
+): boolean {
+	const name = getFunctionName(path);
+	if (!name) return false;
+	if (name === DefaultExportSymbol) {
+		return basename(filename)?.match(/^use[A-Z]/) != null ?? false;
+	}
 
-const fnNameStartsWithCapital = testFunctionName(
-	name => name?.match(/^[A-Z]/) !== null
-);
-const fnNameStartsWithUse = testFunctionName(
-	name => name?.match(/^use[A-Z]/) !== null
-);
+	return name.match(/^use[A-Z]/) != null;
+}
 
 function hasLeadingComment(path: NodePath, comment: RegExp): boolean {
 	const comments = path.node.leadingComments;
@@ -101,9 +187,12 @@ function isOptedIntoSignalTracking(path: NodePath | null): boolean {
 		case "ArrowFunctionExpression":
 		case "FunctionExpression":
 		case "FunctionDeclaration":
+		case "ObjectMethod":
+		case "ObjectExpression":
 		case "VariableDeclarator":
 		case "VariableDeclaration":
 		case "AssignmentExpression":
+		case "CallExpression":
 			return (
 				hasLeadingOptInComment(path) ||
 				isOptedIntoSignalTracking(path.parentPath)
@@ -125,9 +214,12 @@ function isOptedOutOfSignalTracking(path: NodePath | null): boolean {
 		case "ArrowFunctionExpression":
 		case "FunctionExpression":
 		case "FunctionDeclaration":
+		case "ObjectMethod":
+		case "ObjectExpression":
 		case "VariableDeclarator":
 		case "VariableDeclaration":
 		case "AssignmentExpression":
+		case "CallExpression":
 			return (
 				hasLeadingOptOutComment(path) ||
 				isOptedOutOfSignalTracking(path.parentPath)
@@ -142,19 +234,26 @@ function isOptedOutOfSignalTracking(path: NodePath | null): boolean {
 	}
 }
 
-function isComponentFunction(path: NodePath<FunctionLike>): boolean {
+function isComponentFunction(
+	path: NodePath<FunctionLike>,
+	filename: string | undefined
+): boolean {
 	return (
-		fnNameStartsWithCapital(path) && // Function name indicates it's a component
-		getData(path.scope, containsJSX) === true // Function contains JSX
+		getData(path.scope, containsJSX) === true && // Function contains JSX
+		fnNameStartsWithCapital(path, filename) // Function name indicates it's a component
 	);
 }
 
-function isCustomHook(path: NodePath<FunctionLike>): boolean {
-	return fnNameStartsWithUse(path); // Function name indicates it's a hook
+function isCustomHook(
+	path: NodePath<FunctionLike>,
+	filename: string | undefined
+): boolean {
+	return fnNameStartsWithUse(path, filename); // Function name indicates it's a hook
 }
 
 function shouldTransform(
 	path: NodePath<FunctionLike>,
+	filename: string | undefined,
 	options: PluginOptions
 ): boolean {
 	if (getData(path, alreadyTransformed) === true) return false;
@@ -165,14 +264,14 @@ function shouldTransform(
 	if (isOptedIntoSignalTracking(path)) return true;
 
 	if (options.mode === "all") {
-		return isComponentFunction(path);
+		return isComponentFunction(path, filename);
 	}
 
 	if (options.mode == null || options.mode === "auto") {
 		return (
-			(isComponentFunction(path) || isCustomHook(path)) &&
-			getData(path.scope, maybeUsesSignal) === true
-		); // Function appears to use signals;
+			getData(path.scope, maybeUsesSignal) === true && // Function appears to use signals;
+			(isComponentFunction(path, filename) || isCustomHook(path, filename))
+		);
 	}
 
 	return false;
@@ -242,10 +341,11 @@ function transformFunction(
 	t: typeof BabelTypes,
 	options: PluginOptions,
 	path: NodePath<FunctionLike>,
+	filename: string | undefined,
 	state: PluginPass
 ) {
 	let newFunction: FunctionLike;
-	if (isCustomHook(path) || options.experimental?.noTryFinally) {
+	if (isCustomHook(path, filename) || options.experimental?.noTryFinally) {
 		// For custom hooks, we don't need to wrap the function body in a
 		// try/finally block because later code in the function's render body could
 		// read signals and we want to track and associate those signals with this
@@ -369,24 +469,32 @@ export default function signalsTransform(
 				// seeing a function would probably be faster than running an entire
 				// babel pass with plugins on components twice.
 				exit(path, state) {
-					if (shouldTransform(path, options)) {
-						transformFunction(t, options, path, state);
+					if (shouldTransform(path, this.filename, options)) {
+						transformFunction(t, options, path, this.filename, state);
 					}
 				},
 			},
 
 			FunctionExpression: {
 				exit(path, state) {
-					if (shouldTransform(path, options)) {
-						transformFunction(t, options, path, state);
+					if (shouldTransform(path, this.filename, options)) {
+						transformFunction(t, options, path, this.filename, state);
 					}
 				},
 			},
 
 			FunctionDeclaration: {
 				exit(path, state) {
-					if (shouldTransform(path, options)) {
-						transformFunction(t, options, path, state);
+					if (shouldTransform(path, this.filename, options)) {
+						transformFunction(t, options, path, this.filename, state);
+					}
+				},
+			},
+
+			ObjectMethod: {
+				exit(path, state) {
+					if (shouldTransform(path, this.filename, options)) {
+						transformFunction(t, options, path, this.filename, state);
 					}
 				},
 			},