From 64fe47b9e66dacb6617b6f6c41e34049bc38733e Mon Sep 17 00:00:00 2001
From: flavien <flaviendelangle@gmail.com>
Date: Tue, 24 Sep 2024 11:55:12 +0200
Subject: [PATCH] [pickers] Improve typing of the range pickers

---
 .../useDesktopRangePicker.tsx                 | 93 +++++++++----------
 .../hooks/useEnrichedRangePickerFieldProps.ts | 88 +++++++++++-------
 .../useMobileRangePicker.tsx                  | 73 +++++++--------
 3 files changed, 133 insertions(+), 121 deletions(-)

diff --git a/packages/x-date-pickers-pro/src/internals/hooks/useDesktopRangePicker/useDesktopRangePicker.tsx b/packages/x-date-pickers-pro/src/internals/hooks/useDesktopRangePicker/useDesktopRangePicker.tsx
index 55d56d9b6ca58..09d0fc443eba5 100644
--- a/packages/x-date-pickers-pro/src/internals/hooks/useDesktopRangePicker/useDesktopRangePicker.tsx
+++ b/packages/x-date-pickers-pro/src/internals/hooks/useDesktopRangePicker/useDesktopRangePicker.tsx
@@ -11,22 +11,22 @@ import {
   ExportedBaseToolbarProps,
   DateOrTimeViewWithMeridiem,
   ExportedBaseTabsProps,
+  PickersFieldContextValue,
+  PickersFieldProvider,
 } from '@mui/x-date-pickers/internals';
-import {
-  PickerValidDate,
-  FieldRef,
-  BaseSingleInputFieldProps,
-  InferError,
-} from '@mui/x-date-pickers/models';
+import { PickerValidDate, FieldRef, InferError } from '@mui/x-date-pickers/models';
 import {
   DesktopRangePickerAdditionalViewProps,
   UseDesktopRangePickerParams,
   UseDesktopRangePickerProps,
   UseDesktopRangePickerSlotProps,
 } from './useDesktopRangePicker.types';
-import { useEnrichedRangePickerFieldProps } from '../useEnrichedRangePickerFieldProps';
+import {
+  RangePickerPropsForFieldSlot,
+  useEnrichedRangePickerFieldProps,
+} from '../useEnrichedRangePickerFieldProps';
 import { getReleaseInfo } from '../../utils/releaseInfo';
-import { DateRange, BaseMultiInputFieldProps, RangeFieldSection } from '../../../models';
+import { DateRange, RangeFieldSection } from '../../../models';
 import { useRangePosition } from '../useRangePosition';
 
 const releaseInfo = getReleaseInfo();
@@ -138,24 +138,12 @@ export const useDesktopRangePicker = <
   const fieldProps = useSlotProps<
     typeof Field,
     UseDesktopRangePickerSlotProps<TDate, TView, TEnableAccessibleFieldDOMStructure>['field'],
-    | Partial<
-        BaseSingleInputFieldProps<
-          DateRange<TDate>,
-          TDate,
-          RangeFieldSection,
-          TEnableAccessibleFieldDOMStructure,
-          InferError<TExternalProps>
-        >
-      >
-    | Partial<
-        BaseMultiInputFieldProps<
-          DateRange<TDate>,
-          TDate,
-          RangeFieldSection,
-          TEnableAccessibleFieldDOMStructure,
-          InferError<TExternalProps>
-        >
-      >,
+    RangePickerPropsForFieldSlot<
+      boolean,
+      TDate,
+      TEnableAccessibleFieldDOMStructure,
+      InferError<TExternalProps>
+    >,
     TExternalProps
   >({
     elementType: Field,
@@ -222,32 +210,39 @@ export const useDesktopRangePicker = <
   };
   const Layout = slots?.layout ?? PickersLayout;
 
+  const contextValue = React.useMemo<PickersFieldContextValue>(
+    () => ({ onOpen: actions.onOpen }),
+    [actions.onOpen],
+  );
+
   const renderPicker = () => (
-    <LocalizationProvider localeText={localeText}>
-      <Field {...enrichedFieldProps} />
-      <PickersPopper
-        role="tooltip"
-        placement="bottom-start"
-        containerRef={popperRef}
-        anchorEl={anchorRef.current}
-        onBlur={handleBlur}
-        {...actions}
-        open={open}
-        slots={slots}
-        slotProps={slotProps}
-        shouldRestoreFocus={shouldRestoreFocus}
-        reduceAnimations={reduceAnimations}
-      >
-        <Layout
-          {...layoutProps}
-          {...slotProps?.layout}
+    <PickersFieldProvider value={contextValue}>
+      <LocalizationProvider localeText={localeText}>
+        <Field {...enrichedFieldProps} />
+        <PickersPopper
+          role="tooltip"
+          placement="bottom-start"
+          containerRef={popperRef}
+          anchorEl={anchorRef.current}
+          onBlur={handleBlur}
+          {...actions}
+          open={open}
           slots={slots}
-          slotProps={slotPropsForLayout}
+          slotProps={slotProps}
+          shouldRestoreFocus={shouldRestoreFocus}
+          reduceAnimations={reduceAnimations}
         >
-          {renderCurrentView()}
-        </Layout>
-      </PickersPopper>
-    </LocalizationProvider>
+          <Layout
+            {...layoutProps}
+            {...slotProps?.layout}
+            slots={slots}
+            slotProps={slotPropsForLayout}
+          >
+            {renderCurrentView()}
+          </Layout>
+        </PickersPopper>
+      </LocalizationProvider>
+    </PickersFieldProvider>
   );
 
   return {
diff --git a/packages/x-date-pickers-pro/src/internals/hooks/useEnrichedRangePickerFieldProps.ts b/packages/x-date-pickers-pro/src/internals/hooks/useEnrichedRangePickerFieldProps.ts
index 92161a1d50daa..9c0bde6d9d5bb 100644
--- a/packages/x-date-pickers-pro/src/internals/hooks/useEnrichedRangePickerFieldProps.ts
+++ b/packages/x-date-pickers-pro/src/internals/hooks/useEnrichedRangePickerFieldProps.ts
@@ -19,7 +19,6 @@ import {
 } from '@mui/x-date-pickers/hooks';
 import { PickersInputLocaleText } from '@mui/x-date-pickers/locales';
 import {
-  BaseFieldProps,
   onSpaceOrEnter,
   UsePickerResponse,
   WrapperVariant,
@@ -83,24 +82,37 @@ export interface RangePickerFieldSlotProps<
   >;
 }
 
+export type RangePickerPropsForFieldSlot<
+  TIsSingleInput extends boolean,
+  TDate extends PickerValidDate,
+  TEnableAccessibleFieldDOMStructure extends boolean,
+  TError,
+> =
+  | (TIsSingleInput extends true
+      ? BaseSingleInputFieldProps<
+          DateRange<TDate>,
+          TDate,
+          RangeFieldSection,
+          TEnableAccessibleFieldDOMStructure,
+          TError
+        >
+      : never)
+  | (TIsSingleInput extends false
+      ? BaseMultiInputFieldProps<
+          DateRange<TDate>,
+          TDate,
+          RangeFieldSection,
+          TEnableAccessibleFieldDOMStructure,
+          TError
+        >
+      : never);
+
 export interface UseEnrichedRangePickerFieldPropsParams<
+  TIsSingleInput extends boolean,
   TDate extends PickerValidDate,
   TView extends DateOrTimeViewWithMeridiem,
   TEnableAccessibleFieldDOMStructure extends boolean,
   TError,
-  FieldProps extends BaseFieldProps<
-    DateRange<TDate>,
-    TDate,
-    RangeFieldSection,
-    TEnableAccessibleFieldDOMStructure,
-    TError
-  > = BaseFieldProps<
-    DateRange<TDate>,
-    TDate,
-    RangeFieldSection,
-    TEnableAccessibleFieldDOMStructure,
-    TError
-  >,
 > extends Pick<
       UsePickerResponse<DateRange<TDate>, TView, RangeFieldSection, any>,
       'open' | 'actions'
@@ -116,7 +128,12 @@ export interface UseEnrichedRangePickerFieldPropsParams<
   localeText: PickersInputLocaleText<TDate> | undefined;
   pickerSlotProps: RangePickerFieldSlotProps<TDate, TEnableAccessibleFieldDOMStructure> | undefined;
   pickerSlots: RangePickerFieldSlots | undefined;
-  fieldProps: FieldProps;
+  fieldProps: RangePickerPropsForFieldSlot<
+    TIsSingleInput,
+    TDate,
+    TEnableAccessibleFieldDOMStructure,
+    TError
+  >;
   anchorRef?: React.Ref<HTMLDivElement>;
   currentView?: TView | null;
   initialView?: TView;
@@ -151,17 +168,11 @@ const useMultiInputFieldSlotProps = <
   startFieldRef,
   endFieldRef,
 }: UseEnrichedRangePickerFieldPropsParams<
+  false,
   TDate,
   TView,
   TEnableAccessibleFieldDOMStructure,
-  TError,
-  BaseMultiInputFieldProps<
-    DateRange<TDate>,
-    TDate,
-    RangeFieldSection,
-    TEnableAccessibleFieldDOMStructure,
-    TError
-  >
+  TError
 >) => {
   type ReturnType = BaseMultiInputFieldProps<
     DateRange<TDate>,
@@ -336,17 +347,11 @@ const useSingleInputFieldSlotProps = <
   anchorRef,
   currentView,
 }: UseEnrichedRangePickerFieldPropsParams<
+  true,
   TDate,
   TView,
   TEnableAccessibleFieldDOMStructure,
-  TError,
-  BaseSingleInputFieldProps<
-    DateRange<TDate>,
-    TDate,
-    RangeFieldSection,
-    TEnableAccessibleFieldDOMStructure,
-    TError
-  >
+  TError
 >) => {
   type ReturnType = BaseSingleInputFieldProps<
     DateRange<TDate>,
@@ -453,6 +458,7 @@ export const useEnrichedRangePickerFieldProps = <
   TError,
 >(
   params: UseEnrichedRangePickerFieldPropsParams<
+    boolean,
     TDate,
     TView,
     TEnableAccessibleFieldDOMStructure,
@@ -470,9 +476,25 @@ export const useEnrichedRangePickerFieldProps = <
   }
 
   if (params.fieldType === 'multi-input') {
-    return useMultiInputFieldSlotProps(params);
+    return useMultiInputFieldSlotProps(
+      params as unknown as UseEnrichedRangePickerFieldPropsParams<
+        false,
+        TDate,
+        TView,
+        TEnableAccessibleFieldDOMStructure,
+        TError
+      >,
+    );
   }
 
-  return useSingleInputFieldSlotProps(params);
+  return useSingleInputFieldSlotProps(
+    params as UseEnrichedRangePickerFieldPropsParams<
+      true,
+      TDate,
+      TView,
+      TEnableAccessibleFieldDOMStructure,
+      TError
+    >,
+  );
   /* eslint-enable react-hooks/rules-of-hooks */
 };
diff --git a/packages/x-date-pickers-pro/src/internals/hooks/useMobileRangePicker/useMobileRangePicker.tsx b/packages/x-date-pickers-pro/src/internals/hooks/useMobileRangePicker/useMobileRangePicker.tsx
index 518914bdb1b59..d58c24ba6e3e3 100644
--- a/packages/x-date-pickers-pro/src/internals/hooks/useMobileRangePicker/useMobileRangePicker.tsx
+++ b/packages/x-date-pickers-pro/src/internals/hooks/useMobileRangePicker/useMobileRangePicker.tsx
@@ -9,14 +9,11 @@ import {
   ExportedBaseToolbarProps,
   DateOrTimeViewWithMeridiem,
   ExportedBaseTabsProps,
+  PickersFieldContextValue,
+  PickersFieldProvider,
 } from '@mui/x-date-pickers/internals';
 import { usePickersTranslations } from '@mui/x-date-pickers/hooks';
-import {
-  PickerValidDate,
-  FieldRef,
-  BaseSingleInputFieldProps,
-  InferError,
-} from '@mui/x-date-pickers/models';
+import { PickerValidDate, FieldRef, InferError } from '@mui/x-date-pickers/models';
 import useId from '@mui/utils/useId';
 import {
   MobileRangePickerAdditionalViewProps,
@@ -24,9 +21,12 @@ import {
   UseMobileRangePickerProps,
   UseMobileRangePickerSlotProps,
 } from './useMobileRangePicker.types';
-import { useEnrichedRangePickerFieldProps } from '../useEnrichedRangePickerFieldProps';
+import {
+  RangePickerPropsForFieldSlot,
+  useEnrichedRangePickerFieldProps,
+} from '../useEnrichedRangePickerFieldProps';
 import { getReleaseInfo } from '../../utils/releaseInfo';
-import { DateRange, BaseMultiInputFieldProps, RangeFieldSection } from '../../../models';
+import { DateRange, RangeFieldSection } from '../../../models';
 import { useRangePosition } from '../useRangePosition';
 
 const releaseInfo = getReleaseInfo();
@@ -114,24 +114,12 @@ export const useMobileRangePicker = <
   const fieldProps = useSlotProps<
     typeof Field,
     UseMobileRangePickerSlotProps<TDate, TView, TEnableAccessibleFieldDOMStructure>['field'],
-    | Partial<
-        BaseSingleInputFieldProps<
-          DateRange<TDate>,
-          TDate,
-          RangeFieldSection,
-          TEnableAccessibleFieldDOMStructure,
-          InferError<TExternalProps>
-        >
-      >
-    | Partial<
-        BaseMultiInputFieldProps<
-          DateRange<TDate>,
-          TDate,
-          RangeFieldSection,
-          TEnableAccessibleFieldDOMStructure,
-          InferError<TExternalProps>
-        >
-      >,
+    RangePickerPropsForFieldSlot<
+      boolean,
+      TDate,
+      TEnableAccessibleFieldDOMStructure,
+      InferError<TExternalProps>
+    >,
     TExternalProps
   >({
     elementType: Field,
@@ -227,20 +215,27 @@ export const useMobileRangePicker = <
     },
   };
 
+  const contextValue = React.useMemo<PickersFieldContextValue>(
+    () => ({ onOpen: actions.onOpen }),
+    [actions.onOpen],
+  );
+
   const renderPicker = () => (
-    <LocalizationProvider localeText={localeText}>
-      <Field {...enrichedFieldProps} />
-      <PickersModalDialog {...actions} open={open} slots={slots} slotProps={slotProps}>
-        <Layout
-          {...layoutProps}
-          {...slotProps?.layout}
-          slots={slots}
-          slotProps={slotPropsForLayout}
-        >
-          {renderCurrentView()}
-        </Layout>
-      </PickersModalDialog>
-    </LocalizationProvider>
+    <PickersFieldProvider value={contextValue}>
+      <LocalizationProvider localeText={localeText}>
+        <Field {...enrichedFieldProps} />
+        <PickersModalDialog {...actions} open={open} slots={slots} slotProps={slotProps}>
+          <Layout
+            {...layoutProps}
+            {...slotProps?.layout}
+            slots={slots}
+            slotProps={slotPropsForLayout}
+          >
+            {renderCurrentView()}
+          </Layout>
+        </PickersModalDialog>
+      </LocalizationProvider>
+    </PickersFieldProvider>
   );
 
   return {