[pickers] Fix x-date-pickers-pro theme augmentation (#6096)

* Add static pickers to augmentation types

* Fix to correctly pass down `className` to static picker wrapper

* Correctly apply `DateRangePickerDay` style overrides

* Use classes and className in range picker components

packages/x-date-pickers-pro/src/DateRangePicker/DateRangePickerInput.tsx

@@ -1,5 +1,7 @@
 import * as React from 'react';
-import { styled } from '@mui/material/styles';
+import clsx from 'clsx';
+import { styled, useThemeProps } from '@mui/material/styles';
+import { unstable_composeClasses as composeClasses } from '@mui/material';
 import {
   useUtils,
   executeInTheNextEventLoopTick,
@@ -11,8 +13,25 @@ import {
 } from '@mui/x-date-pickers/internals';
 import { CurrentlySelectingRangeEndProps, DateRange } from '../internal/models/dateRange';
 import { DateRangeValidationError } from '../internal/hooks/validation/useDateRangeValidation';
+import {
+  DateRangePickerInputClasses,
+  getDateRangePickerInputUtilityClass,
+} from './dateRangePickerInputClasses';
+
+const useUtilityClasses = (ownerState: DateRangePickerInputProps<any, any>) => {
+  const { classes } = ownerState;
+  const slots = {
+    root: ['root'],
+  };
+
+  return composeClasses(slots, getDateRangePickerInputUtilityClass, classes);
+};
 
-const DateRangePickerInputRoot = styled('div')(({ theme }) => ({
+const DateRangePickerInputRoot = styled('div', {
+  name: 'MuiDateRangePickerInput',
+  slot: 'Root',
+  overridesResolver: (_, styles) => styles.root,
+})(({ theme }) => ({
   display: 'flex',
   alignItems: 'baseline',
   [theme.breakpoints.down('xs')]: {
@@ -48,7 +67,7 @@ export interface ExportedDateRangePickerInputProps<TInputDate, TDate>
   onChange: (date: DateRange<TDate>, keyboardInputValue?: string) => void;
 }
 
-export interface DateRangeInputProps<TInputDate, TDate>
+export interface DateRangePickerInputProps<TInputDate, TDate>
   extends ExportedDateRangePickerInputProps<TInputDate, TDate>,
     Omit<
       DateInputProps<TInputDate, TDate>,
@@ -59,10 +78,11 @@ export interface DateRangeInputProps<TInputDate, TDate>
   endText: React.ReactNode;
   validationError: DateRangeValidationError;
   rawValue: DateRange<TInputDate>;
+  classes?: Partial<DateRangePickerInputClasses>;
 }
 
 type DatePickerInputComponent = <TInputDate, TDate>(
-  props: DateRangeInputProps<TInputDate, TDate> & React.RefAttributes<HTMLDivElement>,
+  props: DateRangePickerInputProps<TInputDate, TDate> & React.RefAttributes<HTMLDivElement>,
 ) => JSX.Element;
 
 /**
@@ -71,7 +91,11 @@ type DatePickerInputComponent = <TInputDate, TDate>(
 export const DateRangePickerInput = React.forwardRef(function DateRangePickerInput<
   TInputDate,
   TDate,
->(props: DateRangeInputProps<TInputDate, TDate>, ref: React.Ref<HTMLDivElement>): JSX.Element {
+>(
+  inProps: DateRangePickerInputProps<TInputDate, TDate>,
+  ref: React.Ref<HTMLDivElement>,
+): JSX.Element {
+  const props = useThemeProps({ props: inProps, name: 'MuiDateRangePickerInput' });
   const {
     currentlySelectingRangeEnd,
     disableOpenPicker,
@@ -88,12 +112,14 @@ export const DateRangePickerInput = React.forwardRef(function DateRangePickerInp
     startText,
     TextFieldProps,
     validationError: [startValidationError, endValidationError],
+    className,
     ...other
   } = props;
 
   const utils = useUtils<TDate>();
   const startRef = React.useRef<HTMLInputElement>(null);
   const endRef = React.useRef<HTMLInputElement>(null);
+  const classes = useUtilityClasses(props);
 
   React.useEffect(() => {
     if (!open) {
@@ -191,7 +217,7 @@ export const DateRangePickerInput = React.forwardRef(function DateRangePickerInp
   });
 
   return (
-    <DateRangePickerInputRoot onBlur={onBlur} ref={ref}>
+    <DateRangePickerInputRoot onBlur={onBlur} ref={ref} className={clsx(classes.root, className)}>
       {renderInput(startInputProps, endInputProps)}
     </DateRangePickerInputRoot>
   );

packages/x-date-pickers-pro/src/DateRangePicker/DateRangePickerToolbar.tsx

@@ -1,7 +1,7 @@
 import * as React from 'react';
 import Typography from '@mui/material/Typography';
-import { styled } from '@mui/material/styles';
-import { generateUtilityClasses } from '@mui/material';
+import { styled, useThemeProps } from '@mui/material/styles';
+import { unstable_composeClasses as composeClasses } from '@mui/material';
 import {
   PickersToolbar,
   PickersToolbarButton,
@@ -11,12 +11,22 @@ import {
   useLocaleText,
 } from '@mui/x-date-pickers/internals';
 import { DateRange, CurrentlySelectingRangeEndProps } from '../internal/models';
+import {
+  DateRangePickerToolbarClasses,
+  getDateRangePickerToolbarUtilityClass,
+} from './dateRangePickerToolbarClasses';
+
+const useUtilityClasses = (ownerState: DateRangePickerToolbarProps<any>) => {
+  const { classes } = ownerState;
+  const slots = {
+    root: ['root'],
+    container: ['container'],
+  };
 
-export const dateRangePickerToolbarClasses = generateUtilityClasses('MuiDateRangePickerToolbar', [
-  'root',
-]);
+  return composeClasses(slots, getDateRangePickerToolbarUtilityClass, classes);
+};
 
-interface DateRangePickerToolbarProps<TDate>
+export interface DateRangePickerToolbarProps<TDate>
   extends CurrentlySelectingRangeEndProps,
     Pick<
       BaseToolbarProps<TDate, DateRange<TDate>>,
@@ -28,12 +38,13 @@ interface DateRangePickerToolbarProps<TDate>
     > {
   startText: React.ReactNode;
   endText: React.ReactNode;
+  classes?: Partial<DateRangePickerToolbarClasses>;
 }
 
 const DateRangePickerToolbarRoot = styled(PickersToolbar, {
   name: 'MuiDateRangePickerToolbar',
   slot: 'Root',
-  overridesResolver: (props, styles) => styles.root,
+  overridesResolver: (_, styles) => styles.root,
 })<{
   ownerState: DateRangePickerToolbarProps<any>;
 }>({
@@ -43,17 +54,22 @@ const DateRangePickerToolbarRoot = styled(PickersToolbar, {
   },
 });
 
-const DateRangePickerToolbarContainer = styled('div')({
+const DateRangePickerToolbarContainer = styled('div', {
+  name: 'MuiDateRangePickerToolbar',
+  slot: 'Container',
+  overridesResolver: (_, styles) => styles.container,
+})({
   display: 'flex',
 });
 
 /**
  * @ignore - internal component.
  */
-export const DateRangePickerToolbar = <TDate extends unknown>(
-  props: DateRangePickerToolbarProps<TDate>,
-) => {
+export const DateRangePickerToolbar = React.forwardRef(function DateRangePickerToolbar<
+  TDate extends unknown,
+>(inProps: DateRangePickerToolbarProps<TDate>, ref: React.Ref<HTMLDivElement>) {
   const utils = useUtils<TDate>();
+  const props = useThemeProps({ props: inProps, name: 'MuiDateRangePickerToolbar' });
 
   const {
     currentlySelectingRangeEnd,
@@ -79,17 +95,19 @@ export const DateRangePickerToolbar = <TDate extends unknown>(
     : endText;
 
   const ownerState = props;
+  const classes = useUtilityClasses(ownerState);
 
   return (
     <DateRangePickerToolbarRoot
       toolbarTitle={toolbarTitle}
       isMobileKeyboardViewOpen={isMobileKeyboardViewOpen}
       toggleMobileKeyboardView={toggleMobileKeyboardView}
       isLandscape={false}
-      className={dateRangePickerToolbarClasses.root}
+      className={classes.root}
       ownerState={ownerState}
+      ref={ref}
     >
-      <DateRangePickerToolbarContainer>
+      <DateRangePickerToolbarContainer className={classes.container}>
         <PickersToolbarButton
           variant={start !== null ? 'h5' : 'h6'}
           value={startDateValue}
@@ -106,4 +124,4 @@ export const DateRangePickerToolbar = <TDate extends unknown>(
       </DateRangePickerToolbarContainer>
     </DateRangePickerToolbarRoot>
   );
-};
+});

packages/x-date-pickers-pro/src/DateRangePicker/DateRangePickerView.tsx

@@ -26,7 +26,7 @@ import {
   DateRangePickerViewMobileSlotsComponent,
   DateRangePickerViewMobileSlotsComponentsProps,
 } from './DateRangePickerViewMobile';
-import { DateRangePickerInput, DateRangeInputProps } from './DateRangePickerInput';
+import { DateRangePickerInput, DateRangePickerInputProps } from './DateRangePickerInput';
 import {
   DateRangePickerViewDesktop,
   ExportedDesktopDateRangeCalendarProps,
@@ -94,7 +94,7 @@ interface DateRangePickerViewProps<TInputDate, TDate>
   open: boolean;
   startText: React.ReactNode;
   endText: React.ReactNode;
-  DateInputProps: DateRangeInputProps<TInputDate, TDate>;
+  DateInputProps: DateRangePickerInputProps<TInputDate, TDate>;
 }
 
 type DateRangePickerViewComponent = (<TInputDate, TDate = TInputDate>(

packages/x-date-pickers-pro/src/DateRangePicker/DateRangePickerViewDesktop.tsx

@@ -1,5 +1,7 @@
 import * as React from 'react';
-import { styled } from '@mui/material/styles';
+import clsx from 'clsx';
+import { styled, useThemeProps } from '@mui/material/styles';
+import { unstable_composeClasses as composeClasses } from '@mui/material';
 import {
   useDefaultDates,
   useUtils,
@@ -19,6 +21,20 @@ import { DateRange } from '../internal/models';
 import { DateRangePickerDay, DateRangePickerDayProps } from '../DateRangePickerDay';
 import { isWithinRange, isStartOfRange, isEndOfRange } from '../internal/utils/date-utils';
 import { doNothing } from '../internal/utils/utils';
+import {
+  DateRangePickerViewDesktopClasses,
+  getDateRangePickerViewDesktopUtilityClass,
+} from './dateRangePickerViewDesktopClasses';
+
+const useUtilityClasses = (ownerState: DesktopDateRangeCalendarProps<any>) => {
+  const { classes } = ownerState;
+  const slots = {
+    root: ['root'],
+    container: ['container'],
+  };
+
+  return composeClasses(slots, getDateRangePickerViewDesktopUtilityClass, classes);
+};
 
 export interface ExportedDesktopDateRangeCalendarProps<TDate> {
   /**
@@ -37,7 +53,7 @@ export interface ExportedDesktopDateRangeCalendarProps<TDate> {
   renderDay?: (day: TDate, dateRangePickerDayProps: DateRangePickerDayProps<TDate>) => JSX.Element;
 }
 
-interface DesktopDateRangeCalendarProps<TDate>
+export interface DesktopDateRangeCalendarProps<TDate>
   extends ExportedDesktopDateRangeCalendarProps<TDate>,
     Omit<DayPickerProps<TDate>, 'selectedDays' | 'renderDay' | 'onFocusedDayChange'>,
     DayValidationProps<TDate>,
@@ -46,14 +62,23 @@ interface DesktopDateRangeCalendarProps<TDate>
   parsedValue: DateRange<TDate>;
   changeMonth: (date: TDate) => void;
   currentlySelectingRangeEnd: 'start' | 'end';
+  classes?: Partial<DateRangePickerViewDesktopClasses>;
 }
 
-const DateRangePickerViewDesktopRoot = styled('div')({
+const DateRangePickerViewDesktopRoot = styled('div', {
+  name: 'MuiDateRangePickerViewDesktop',
+  slot: 'Root',
+  overridesResolver: (_, styles) => styles.root,
+})({
   display: 'flex',
   flexDirection: 'row',
 });
 
-const DateRangePickerViewDesktopContainer = styled('div')(({ theme }) => ({
+const DateRangePickerViewDesktopContainer = styled('div', {
+  name: 'MuiDateRangePickerViewDesktop',
+  slot: 'Container',
+  overridesResolver: (_, styles) => styles.container,
+})(({ theme }) => ({
   '&:not(:last-of-type)': {
     borderRight: `2px solid ${theme.palette.divider}`,
   },
@@ -96,7 +121,8 @@ const deprecatedPropsWarning = buildDeprecatedPropsWarning(
 /**
  * @ignore - internal component.
  */
-export function DateRangePickerViewDesktop<TDate>(props: DesktopDateRangeCalendarProps<TDate>) {
+export function DateRangePickerViewDesktop<TDate>(inProps: DesktopDateRangeCalendarProps<TDate>) {
+  const props = useThemeProps({ props: inProps, name: 'MuiDateRangePickerViewDesktop' });
   const {
     calendars,
     changeMonth,
@@ -113,6 +139,7 @@ export function DateRangePickerViewDesktop<TDate>(props: DesktopDateRangeCalenda
     onSelectedDaysChange,
     renderDay = (_, dateRangeProps) => <DateRangePickerDay {...dateRangeProps} />,
     rightArrowButtonText: rightArrowButtonTextProp,
+    className,
     ...other
   } = props;
 
@@ -127,6 +154,7 @@ export function DateRangePickerViewDesktop<TDate>(props: DesktopDateRangeCalenda
   const rightArrowButtonText = rightArrowButtonTextProp ?? localeText.nextMonth;
 
   const utils = useUtils<TDate>();
+  const classes = useUtilityClasses(props);
   const defaultDates = useDefaultDates<TDate>();
   const minDate = minDateProp ?? defaultDates.minDate;
   const maxDate = maxDateProp ?? defaultDates.maxDate;
@@ -175,12 +203,12 @@ export function DateRangePickerViewDesktop<TDate>(props: DesktopDateRangeCalenda
   }, [changeMonth, currentMonth, utils]);
 
   return (
-    <DateRangePickerViewDesktopRoot>
+    <DateRangePickerViewDesktopRoot className={clsx(className, classes.root)}>
       {getCalendarsArray(calendars).map((_, index) => {
         const monthOnIteration = utils.setMonth(currentMonth, utils.getMonth(currentMonth) + index);
 
         return (
-          <DateRangePickerViewDesktopContainer key={index}>
+          <DateRangePickerViewDesktopContainer key={index} className={classes.container}>
             <DateRangePickerViewDesktopArrowSwitcher
               onLeftClick={selectPreviousMonth}
               onRightClick={selectNextMonth}

packages/x-date-pickers-pro/src/DateRangePicker/dateRangePickerInputClasses.ts

@@ -0,0 +1,17 @@
+import { generateUtilityClass, generateUtilityClasses } from '@mui/material';
+
+export interface DateRangePickerInputClasses {
+  /** Styles applied to the root element. */
+  root: string;
+}
+
+export type DateRangePickerInputClassKey = keyof DateRangePickerInputClasses;
+
+export function getDateRangePickerInputUtilityClass(slot: string) {
+  return generateUtilityClass('MuiDateRangePickerInput', slot);
+}
+
+export const dateRangePickerInputClasses: DateRangePickerInputClasses = generateUtilityClasses(
+  'MuiDateRangePickerInput',
+  ['root'],
+);

packages/x-date-pickers-pro/src/DateRangePicker/dateRangePickerToolbarClasses.ts

@@ -0,0 +1,19 @@
+import { generateUtilityClass, generateUtilityClasses } from '@mui/material';
+
+export interface DateRangePickerToolbarClasses {
+  /** Styles applied to the root element. */
+  root: string;
+  /** Styles applied to the container element. */
+  container: string;
+}
+
+export type DateRangePickerToolbarClassKey = keyof DateRangePickerToolbarClasses;
+
+export function getDateRangePickerToolbarUtilityClass(slot: string) {
+  return generateUtilityClass('MuiDateRangePickerToolbar', slot);
+}
+
+export const dateRangePickerToolbarClasses: DateRangePickerToolbarClasses = generateUtilityClasses(
+  'MuiDateRangePickerToolbar',
+  ['root', 'container'],
+);

packages/x-date-pickers-pro/src/DateRangePicker/dateRangePickerViewDesktopClasses.ts

@@ -0,0 +1,17 @@
+import { generateUtilityClass, generateUtilityClasses } from '@mui/material';
+
+export interface DateRangePickerViewDesktopClasses {
+  /** Styles applied to the root element. */
+  root: string;
+  /** Styles applied to the container element. */
+  container: string;
+}
+
+export type DateRangePickerViewDesktopClassKey = keyof DateRangePickerViewDesktopClasses;
+
+export function getDateRangePickerViewDesktopUtilityClass(slot: string) {
+  return generateUtilityClass('MuiDateRangePickerViewDesktop', slot);
+}
+
+export const dateRangePickerViewDesktopClasses: DateRangePickerViewDesktopClasses =
+  generateUtilityClasses('MuiDateRangePickerViewDesktop', ['root', 'container']);