[TypeScript] Preserve the generic signature of getEntityRecord and getEntityRecords through currying

packages/core-data/src/selectors.ts

@@ -221,6 +221,55 @@ export function getEntityConfig(
 	return find( state.entities.config, { kind, name } );
 }
 
+/**
+ * GetEntityRecord is declared as a *callable interface* with
+ * two signatures to work around the fact that TypeScript doesn't
+ * allow currying generic functions:
+ *
+ * ```ts
+ * 		type CurriedState = F extends ( state: any, ...args: infer P ) => infer R
+ * 			? ( ...args: P ) => R
+ * 			: F;
+ * 		type Selector = <K extends string | number>(
+ *         state: any,
+ *         kind: K,
+ *         key: K extends string ? 'string value' : false
+ *    ) => K;
+ * 		type BadlyInferredSignature = CurriedState< Selector >
+ *    // BadlyInferredSignature evaluates to:
+ *    // (kind: string number, key: false | "string value") => string number
+ * ```
+ *
+ * The signature without the state parameter shipped as CurriedSignature
+ * is used in the return value of `select( coreStore )`.
+ *
+ * See https://github.com/WordPress/gutenberg/pull/41578 for more details.
+ */
+export interface GetEntityRecord {
+	<
+		EntityRecord extends
+			| ET.EntityRecord< any >
+			| Partial< ET.EntityRecord< any > >
+	>(
+		state: State,
+		kind: string,
+		name: string,
+		key: EntityRecordKey,
+		query?: GetRecordsHttpQuery
+	): EntityRecord | undefined;
+
+	CurriedSignature: <
+		EntityRecord extends
+			| ET.EntityRecord< any >
+			| Partial< ET.EntityRecord< any > >
+	>(
+		kind: string,
+		name: string,
+		key: EntityRecordKey,
+		query?: GetRecordsHttpQuery
+	) => EntityRecord | undefined;
+}
+
 /**
  * Returns the Entity's record object by key. Returns `null` if the value is not
  * yet received, undefined if the value entity is known to not exist, or the
@@ -236,7 +285,7 @@ export function getEntityConfig(
  * @return Record.
  */
 export const getEntityRecord = createSelector(
-	<
+	( <
 		EntityRecord extends
 			| ET.EntityRecord< any >
 			| Partial< ET.EntityRecord< any > >
@@ -279,7 +328,7 @@ export const getEntityRecord = createSelector(
 		}
 
 		return item;
-	},
+	} ) as GetEntityRecord,
 	( state: State, kind, name, recordId, query ) => {
 		const context = query?.context ?? 'default';
 		return [
@@ -301,7 +350,7 @@ export const getEntityRecord = createSelector(
 			] ),
 		];
 	}
-);
+) as GetEntityRecord;
 
 /**
  * Returns the Entity's record object by key. Doesn't trigger a resolver nor requests the entity records from the API if the entity record isn't available in the local state.
@@ -414,6 +463,37 @@ export function hasEntityRecords(
 	return Array.isArray( getEntityRecords( state, kind, name, query ) );
 }
 
+/**
+ * GetEntityRecord is declared as a *callable interface* with
+ * two signatures to work around the fact that TypeScript doesn't
+ * allow currying generic functions.
+ *
+ * @see GetEntityRecord
+ * @see https://github.com/WordPress/gutenberg/pull/41578
+ */
+export interface GetEntityRecords {
+	<
+		EntityRecord extends
+			| ET.EntityRecord< any >
+			| Partial< ET.EntityRecord< any > >
+	>(
+		state: State,
+		kind: string,
+		name: string,
+		query?: GetRecordsHttpQuery
+	): EntityRecord[] | null;
+
+	CurriedSignature: <
+		EntityRecord extends
+			| ET.EntityRecord< any >
+			| Partial< ET.EntityRecord< any > >
+	>(
+		kind: string,
+		name: string,
+		query?: GetRecordsHttpQuery
+	) => EntityRecord[] | null;
+}
+
 /**
  * Returns the Entity's records.
  *
@@ -425,15 +505,15 @@ export function hasEntityRecords(
  *
  * @return Records.
  */
-export const getEntityRecords = <
+export const getEntityRecords = ( <
 	EntityRecord extends
 		| ET.EntityRecord< any >
 		| Partial< ET.EntityRecord< any > >
 >(
 	state: State,
 	kind: string,
 	name: string,
-	query?: GetRecordsHttpQuery
+	query: GetRecordsHttpQuery
 ): EntityRecord[] | null => {
 	// Queried data state is prepopulated for all known entities. If this is not
 	// assigned for the given parameters, then it is known to not exist.
@@ -446,7 +526,7 @@ export const getEntityRecords = <
 		return null;
 	}
 	return getQueriedItems( queriedState, query );
-};
+} ) as GetEntityRecords;
 
 type DirtyEntityRecord = {
 	title: string;

packages/data/src/types.ts

@@ -77,22 +77,75 @@ export type CurriedSelectorsOf< S > = S extends StoreDescriptor<
 	: never;
 
 /**
- * Removes the first argument from a function
+ * Removes the first argument from a function.
  *
- * This is designed to remove the `state` parameter from
+ * By default, it removes the `state` parameter from
  * registered selectors since that argument is supplied
  * by the editor when calling `select(…)`.
  *
  * For functions with no arguments, which some selectors
  * are free to define, returns the original function.
+ *
+ * It is possible to manually provide a custom curried signature
+ * and avoid the automatic inference. When the
+ * F generic argument passed to this helper extends the
+ * SelectorWithCustomCurrySignature type, the F['CurriedSignature']
+ * property is used verbatim.
+ *
+ * This is useful because TypeScript does not correctly remove
+ * arguments from complex function signatures constrained by
+ * interdependent generic parameters.
+ * For more context, see https://github.com/WordPress/gutenberg/pull/41578
  */
-export type CurriedState< F > = F extends (
-	state: any,
-	...args: infer P
-) => infer R
+type CurriedState< F > = F extends SelectorWithCustomCurrySignature
+	? F[ 'CurriedSignature' ]
+	: F extends ( state: any, ...args: infer P ) => infer R
 	? ( ...args: P ) => R
 	: F;
 
+/**
+ * Utility to manually specify curried selector signatures.
+ *
+ * It comes handy when TypeScript can't automatically produce the
+ * correct curried function signature. For example:
+ *
+ * ```ts
+ * type BadlyInferredSignature = CurriedState<
+ *     <K extends string | number>(
+ *         state: any,
+ *         kind: K,
+ *         key: K extends string ? 'one value' : false
+ *     ) => K
+ * >
+ * // BadlyInferredSignature evaluates to:
+ * // (kind: string number, key: false "one value") => string number
+ * ```
+ *
+ * With SelectorWithCustomCurrySignature, we can provide a custom
+ * signature and avoid relying on TypeScript inference:
+ * ```ts
+ * interface MySelectorSignature extends SelectorWithCustomCurrySignature {
+ *     <K extends string | number>(
+ *         state: any,
+ *         kind: K,
+ *         key: K extends string ? 'one value' : false
+ *     ): K;
+ *
+ *     CurriedSignature: <K extends string | number>(
+ *         kind: K,
+ *         key: K extends string ? 'one value' : false
+ *     ): K;
+ * }
+ * type CorrectlyInferredSignature = CurriedState<MySelectorSignature>
+ * // <K extends string | number>(kind: K, key: K extends string ? 'one value' : false): K;
+ *
+ * For even more context, see https://github.com/WordPress/gutenberg/pull/41578
+ * ```
+ */
+export interface SelectorWithCustomCurrySignature {
+	CurriedSignature: Function;
+}
+
 export interface DataRegistry {
 	register: ( store: StoreDescriptor< any > ) => void;
 }

packages/docgen/lib/get-type-annotation.js

@@ -405,18 +405,24 @@ function unwrapWrappedSelectors( token ) {
 		return token;
 	}
 
+	if ( babelTypes.isTSAsExpression( token ) ) {
+		// ( ( state, queryId ) => state.queries[ queryId ] ) as any;
+		// \------------------------------------------------/ CallExpression.expression
+		return unwrapWrappedSelectors( token.expression );
+	}
+
 	if ( babelTypes.isCallExpression( token ) ) {
 		// createSelector( ( state, queryId ) => state.queries[ queryId ] );
 		//                 \--------------------------------------------/ CallExpression.arguments[0]
 		if ( token.callee.name === 'createSelector' ) {
-			return token.arguments[ 0 ];
+			return unwrapWrappedSelectors( token.arguments[ 0 ] );
 		}
 
 		// createRegistrySelector( ( selector ) => ( state, queryId ) => select( 'core/queries' ).get( queryId ) );
 		//                                         \-----------------------------------------------------------/ CallExpression.arguments[0].body
 		//                         \---------------------------------------------------------------------------/ CallExpression.arguments[0]
 		if ( token.callee.name === 'createRegistrySelector' ) {
-			return token.arguments[ 0 ].body;
+			return unwrapWrappedSelectors( token.arguments[ 0 ].body );
 		}
 	}
 }

packages/docgen/test/get-type-annotation.js

@@ -394,46 +394,53 @@ describe( 'Type annotations', () => {
 	} );
 
 	describe( 'statically-wrapped function exceptions', () => {
-		it( 'should find types for inner function with `createSelector`', () => {
-			const { tokens } = engine(
-				'test.ts',
-				`/**
-					 * Returns the number of things
-					 *
-					 * @param state - stores all the things
-					 */
-					export const getCount = createSelector( ( state: string[] ) => state.length );
-					`
+		const getStateArgType = ( code ) => {
+			const { tokens } = engine( 'test.ts', code );
+			return getTypeAnnotation(
+				{ tag: 'param', name: 'state' },
+				tokens[ 0 ],
+				0
 			);
+		};
 
-			expect(
-				getTypeAnnotation(
-					{ tag: 'param', name: 'state' },
-					tokens[ 0 ],
-					0
-				)
-			).toBe( 'string[]' );
+		const docString = `/**
+			 * Returns the number of things
+			 *
+			 * @param state - stores all the things
+			 */`;
+		it( 'should find types for a typecasted function', () => {
+			const code = `${ docString }
+			 export const getCount = ( state: string[] ) => state.length;
+			 `;
+			expect( getStateArgType( code ) ).toBe( 'string[]' );
 		} );
 
-		it( 'should find types for inner function with `createRegistrySelector`', () => {
-			const { tokens } = engine(
-				'test.ts',
-				`/**
-					 * Returns the number of things
-					 *
-					 * @param state - stores all the things
-					 */
-					export const getCount = createRegistrySelector( ( select ) => ( state: number ) => state );
-					`
-			);
+		it( 'should find types for a doubly typecasted function', () => {
+			const code = `${ docString }
+				 export const getCount = ( ( state: string[] ) => state.length ) as any as any;
+			 `;
+			expect( getStateArgType( code ) ).toBe( 'string[]' );
+		} );
 
-			expect(
-				getTypeAnnotation(
-					{ tag: 'param', name: 'state' },
-					tokens[ 0 ],
-					0
-				)
-			).toBe( 'number' );
+		it( 'should find types for inner function with `createSelector`', () => {
+			const code = `${ docString }
+				 export const getCount = createSelector( ( state: string[] ) => state.length );
+			 `;
+			expect( getStateArgType( code ) ).toBe( 'string[]' );
+		} );
+
+		it( 'should find types for inner typecasted function with `createSelector`', () => {
+			const code = `${ docString }
+				 export const getCount = createSelector( (( state: string[] ) => state.length) as any );
+			 `;
+			expect( getStateArgType( code ) ).toBe( 'string[]' );
+		} );
+
+		it( 'should find types for inner function with `createRegistrySelector`', () => {
+			const code = `${ docString }
+				 export const getCount = createRegistrySelector( ( select ) => ( state: number ) => state );
+			 `;
+			expect( getStateArgType( code ) ).toBe( 'number' );
 		} );
 	} );
 } );