Components: Add contextConnectWithoutRef() to bypass ref forwarding (#43611)

* Remove unused memo option

* Clean up ts-expect-errors

* Force check two params

* Check contextConnect correctness with TS

* Add tests

* Add more TS tests and fixup types

* Add comment

* Move conditional upstream

Author: mirka

packages/components/src/ui/context/context-connect.ts

@@ -6,7 +6,7 @@ import type { ForwardedRef, ReactChild, ReactNode } from 'react';
 /**
  * WordPress dependencies
  */
-import { forwardRef, memo } from '@wordpress/element';
+import { forwardRef } from '@wordpress/element';
 import warn from '@wordpress/warning';
 
 /**
@@ -16,42 +16,74 @@ import { CONNECT_STATIC_NAMESPACE } from './constants';
 import { getStyledClassNameFromKey } from './get-styled-class-name-from-key';
 import type { WordPressComponentFromProps } from '.';
 
+type AcceptsTwoArgs<
+	F extends ( ...args: any ) => any,
+	ErrorMessage = never
+> = Parameters< F >[ 'length' ] extends 2 ? {} : ErrorMessage;
+
 type ContextConnectOptions = {
-	/** Defaults to `false`. */
-	memo?: boolean;
+	forwardsRef?: boolean;
 };
 
 /**
  * Forwards ref (React.ForwardRef) and "Connects" (or registers) a component
  * within the Context system under a specified namespace.
  *
- * This is an (experimental) evolution of the initial connect() HOC.
- * The hope is that we can improve render performance by removing functional
- * component wrappers.
+ * @param  Component The component to register into the Context system.
+ * @param  namespace The namespace to register the component under.
+ * @return The connected WordPressComponent
+ */
+export function contextConnect<
+	C extends ( props: any, ref: ForwardedRef< any > ) => JSX.Element | null
+>(
+	Component: C &
+		AcceptsTwoArgs<
+			C,
+			'Warning: Your component function does not take a ref as the second argument. Did you mean to use `contextConnectWithoutRef`?'
+		>,
+	namespace: string
+) {
+	return _contextConnect( Component, namespace, { forwardsRef: true } );
+}
+
+/**
+ * "Connects" (or registers) a component within the Context system under a specified namespace.
+ * Does not forward a ref.
  *
  * @param  Component The component to register into the Context system.
  * @param  namespace The namespace to register the component under.
- * @param  options
  * @return The connected WordPressComponent
  */
-export function contextConnect< P >(
-	Component: ( props: P, ref: ForwardedRef< any > ) => JSX.Element | null,
-	namespace: string,
-	options: ContextConnectOptions = {}
-): WordPressComponentFromProps< P > {
-	const { memo: memoProp = false } = options;
+export function contextConnectWithoutRef< P >(
+	Component: ( props: P ) => JSX.Element | null,
+	namespace: string
+) {
+	return _contextConnect( Component, namespace );
+}
 
-	let WrappedComponent = forwardRef( Component );
-	if ( memoProp ) {
-		// @ts-ignore
-		WrappedComponent = memo( WrappedComponent );
-	}
+// This is an (experimental) evolution of the initial connect() HOC.
+// The hope is that we can improve render performance by removing functional
+// component wrappers.
+function _contextConnect<
+	C extends ( props: any, ref: ForwardedRef< any > ) => JSX.Element | null,
+	O extends ContextConnectOptions
+>(
+	Component: C,
+	namespace: string,
+	options?: O
+): WordPressComponentFromProps<
+	Parameters< C >[ 0 ],
+	O[ 'forwardsRef' ] extends true ? true : false
+> {
+	const WrappedComponent = options?.forwardsRef
+		? forwardRef< any, Parameters< C >[ 0 ] >( Component )
+		: Component;
 
 	if ( typeof namespace === 'undefined' ) {
 		warn( 'contextConnect: Please provide a namespace' );
 	}
 
-	// @ts-ignore internal property
+	// @ts-expect-error internal property
 	let mergedNamespace = WrappedComponent[ CONNECT_STATIC_NAMESPACE ] || [
 		namespace,
 	];
@@ -66,18 +98,13 @@ export function contextConnect< P >(
 		mergedNamespace = [ ...mergedNamespace, namespace ];
 	}
 
-	WrappedComponent.displayName = namespace;
-
-	// @ts-ignore internal property
-	WrappedComponent[ CONNECT_STATIC_NAMESPACE ] = [
-		...new Set( mergedNamespace ),
-	];
-
-	// @ts-ignore WordPressComponent property
-	WrappedComponent.selector = `.${ getStyledClassNameFromKey( namespace ) }`;
-
-	// @ts-ignore
-	return WrappedComponent;
+	// @ts-expect-error We can't rely on inferred types here because of the
+	// `as` prop polymorphism we're handling in https://github.com/WordPress/gutenberg/blob/9620bae6fef4fde7cc2b7833f416e240207cda29/packages/components/src/ui/context/wordpress-component.ts#L32-L33
+	return Object.assign( WrappedComponent, {
+		[ CONNECT_STATIC_NAMESPACE ]: [ ...new Set( mergedNamespace ) ],
+		displayName: namespace,
+		selector: `.${ getStyledClassNameFromKey( namespace ) }`,
+	} );
 }
 
 /**

packages/components/src/ui/context/test/context-connect.tsx

@@ -0,0 +1,55 @@
+/**
+ * External dependencies
+ */
+import type { ForwardedRef } from 'react';
+
+/**
+ * Internal dependencies
+ */
+import { contextConnect, contextConnectWithoutRef } from '../context-connect';
+import type { WordPressComponentProps } from '../wordpress-component';
+
+// Static TypeScript tests
+describe( 'ref forwarding', () => {
+	const ComponentWithRef = (
+		props: WordPressComponentProps< {}, 'div' >,
+		ref: ForwardedRef< any >
+	) => <div { ...props } ref={ ref } />;
+	const ComponentWithoutRef = (
+		props: WordPressComponentProps< {}, 'div' >
+	) => <div { ...props } />;
+
+	it( 'should not trigger a TS error if components are passed to the correct connect* functions', () => {
+		contextConnect( ComponentWithRef, 'Foo' );
+		contextConnectWithoutRef( ComponentWithoutRef, 'Foo' );
+	} );
+
+	it( 'should trigger a TS error if components are passed to the wrong connect* functions', () => {
+		// Wrapped in a thunk because React.forwardRef() will throw a console warning if this is executed
+		// eslint-disable-next-line no-unused-expressions
+		() => {
+			// @ts-expect-error This should error
+			contextConnect( ComponentWithoutRef, 'Foo' );
+		};
+
+		// @ts-expect-error This should error
+		contextConnectWithoutRef( ComponentWithRef, 'Foo' );
+	} );
+
+	it( 'should result in a component with the correct prop types', () => {
+		const AcceptsRef = contextConnect( ComponentWithRef, 'Foo' );
+
+		<AcceptsRef ref={ null } />;
+
+		// @ts-expect-error An unaccepted prop should trigger an error
+		<AcceptsRef foo={ null } />;
+
+		const NoRef = contextConnectWithoutRef( ComponentWithoutRef, 'Foo' );
+
+		// @ts-expect-error The ref prop should trigger an error
+		<NoRef ref={ null } />;
+
+		// @ts-expect-error An unaccepted prop should trigger an error
+		<NoRef foo={ null } />;
+	} );
+} );