From 4b7ab27c825bb21eddb5912b7d3392e82db69eda Mon Sep 17 00:00:00 2001
From: ramon <ramonjd@gmail.com>
Date: Tue, 28 May 2024 16:11:32 +1000
Subject: [PATCH] Testing passing an allowedUnitValues prop to UnitControl,
 which intends (eventually) to act in some way like `useCustomUnits` hook, but
 for now just filters the default units and preserves the order.

---
 .../global-styles/background-panel.js         |  1 +
 .../components/src/unit-control/index.tsx     | 26 ++++++++++++++++---
 packages/components/src/unit-control/types.ts |  4 +++
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/packages/block-editor/src/components/global-styles/background-panel.js b/packages/block-editor/src/components/global-styles/background-panel.js
index 18ecc26b002d0e..5e05fe1344a91e 100644
--- a/packages/block-editor/src/components/global-styles/background-panel.js
+++ b/packages/block-editor/src/components/global-styles/background-panel.js
@@ -599,6 +599,7 @@ function BackgroundSizeControls( {
 					size="__unstable-large"
 					__unstableInputWidth="100px"
 					min={ 0 }
+					allowedUnitValues={ [ '%', 'px', 'em', 'rem', 'vw', 'vh' ] }
 					placeholder={ __( 'Auto' ) }
 					disabled={
 						currentValueForToggle !== 'auto' ||
diff --git a/packages/components/src/unit-control/index.tsx b/packages/components/src/unit-control/index.tsx
index 2dd08cc155225f..775e2440dfc140 100644
--- a/packages/components/src/unit-control/index.tsx
+++ b/packages/components/src/unit-control/index.tsx
@@ -25,7 +25,11 @@ import {
 } from './utils';
 import { useControlledState } from '../utils/hooks';
 import { escapeRegExp } from '../utils/strings';
-import type { UnitControlProps, UnitControlOnChangeCallback } from './types';
+import type {
+	UnitControlProps,
+	UnitControlOnChangeCallback,
+	WPUnitControlUnit,
+} from './types';
 import { useDeprecated36pxDefaultSizeProp } from '../utils/use-deprecated-props';
 
 function UnforwardedUnitControl(
@@ -53,6 +57,7 @@ function UnforwardedUnitControl(
 		size = 'default',
 		unit: unitProp,
 		units: unitsProp = CSS_UNITS,
+		allowedUnitValues,
 		value: valueProp,
 		onFocus: onFocusProp,
 		...props
@@ -71,11 +76,26 @@ function UnforwardedUnitControl(
 	// still passes `null` as a `value`.
 	const nonNullValueProp = valueProp ?? undefined;
 	const [ units, reFirstCharacterOfUnits ] = useMemo( () => {
-		const list = getUnitsWithCurrentUnit(
+		let list = getUnitsWithCurrentUnit(
 			nonNullValueProp,
 			unitProp,
 			unitsProp
 		);
+		// If allowedUnitValues is provided, we need to filter the list of units.
+		if ( allowedUnitValues?.length ) {
+			const filteredUnits: WPUnitControlUnit[] = [];
+			allowedUnitValues.forEach( ( allowedUnit ) => {
+				const filteredUnit = list.find(
+					( unit ) => unit.value === allowedUnit
+				);
+				if ( filteredUnit ) {
+					filteredUnits.push( filteredUnit );
+				}
+			} );
+			if ( filteredUnits.length ) {
+				list = filteredUnits;
+			}
+		}
 		const [ { value: firstUnitValue = '' } = {}, ...rest ] = list;
 		const firstCharacters = rest.reduce(
 			( carry, { value } ) => {
@@ -87,7 +107,7 @@ function UnforwardedUnitControl(
 			escapeRegExp( firstUnitValue.substring( 0, 1 ) )
 		);
 		return [ list, new RegExp( `^(?:${ firstCharacters })$`, 'i' ) ];
-	}, [ nonNullValueProp, unitProp, unitsProp ] );
+	}, [ nonNullValueProp, unitProp, unitsProp, allowedUnitValues ] );
 	const [ parsedQuantity, parsedUnit ] = getParsedQuantityAndUnit(
 		nonNullValueProp,
 		unitProp,
diff --git a/packages/components/src/unit-control/types.ts b/packages/components/src/unit-control/types.ts
index 9164502668a2b0..326bb4df4b6ec3 100644
--- a/packages/components/src/unit-control/types.ts
+++ b/packages/components/src/unit-control/types.ts
@@ -107,4 +107,8 @@ export type UnitControlProps = Pick< InputControlProps, 'size' > &
 		 * Callback when either the quantity or the unit inputs gains focus.
 		 */
 		onFocus?: FocusEventHandler< HTMLInputElement | HTMLSelectElement >;
+		/**
+		 * A list of allowed unit values. If provided, only these values will be selectable.
+		 */
+		allowedUnitValues?: string[];
 	};