Enhance Dialog so that it allows providing a TransitionComponent

src/components/feedback/Dialog.tsx

@@ -1,6 +1,15 @@
 import classnames from 'classnames';
 import type { RefObject } from 'preact';
-import { useEffect, useLayoutEffect, useRef } from 'preact/hooks';
+import type { JSX } from 'preact';
+import { Fragment } from 'preact';
+import {
+  useCallback,
+  useEffect,
+  useLayoutEffect,
+  useMemo,
+  useRef,
+  useState,
+} from 'preact/hooks';
 
 import { useClickAway } from '../../hooks/use-click-away';
 import { useFocusAway } from '../../hooks/use-focus-away';
@@ -12,6 +21,11 @@ import { downcastRef } from '../../util/typing';
 import Panel from '../layout/Panel';
 import type { PanelProps } from '../layout/Panel';
 
+type TransitionComponent = JSX.ElementType<{
+  visible: boolean;
+  onTransitionEnd?: () => void;
+}>;
+
 type ComponentProps = {
   closeOnClickAway?: boolean;
   closeOnEscape?: boolean;
@@ -33,6 +47,15 @@ type ComponentProps = {
    * Restore focus to previously-focused element when unmounted/closed
    */
   restoreFocus?: boolean;
+
+  /**
+   * Providing this has the next implications:
+   * - The component will be used to wrap the Dialog contents.
+   * - If initialFocus === 'auto', the Dialog will be focused once the open
+   *   transition has finished.
+   * - onClose will be invoked after the close transition has finished.
+   */
+  transitionComponent?: TransitionComponent;
 };
 
 // This component forwards a number of props on to `Panel` but always sets the
@@ -53,6 +76,7 @@ const DialogNext = function Dialog({
   children,
   initialFocus = 'auto',
   restoreFocus = false,
+  transitionComponent: TransitionComponent,
 
   classes,
   elementRef,
@@ -68,31 +92,22 @@ const DialogNext = function Dialog({
   ...htmlAttributes
 }: DialogProps) {
   const modalRef = useSyncedRef(elementRef);
-  const closeHandler = onClose ?? noop;
   const restoreFocusEl = useRef<HTMLElement | null>(
     document.activeElement as HTMLElement | null
   );
-
-  useClickAway(modalRef, closeHandler, {
-    enabled: closeOnClickAway,
-  });
-
-  useKeyPress(['Escape'], closeHandler, {
-    enabled: closeOnEscape,
-  });
-
-  useFocusAway(modalRef, closeHandler, {
-    enabled: closeOnFocusAway,
-  });
-
-  const dialogDescriptionId = useUniqueId('dialog-description');
-
-  useEffect(() => {
+  const [visible, setVisible] = useState(false);
+  // If a TransitionComponent was provided, closing the Dialog should just set
+  // it to not visible. The TransitionComponent will take care of actually
+  // closing the dialog once transition has finished
+  const closeHandler = TransitionComponent
+    ? () => setVisible(false)
+    : onClose ?? noop;
+  const setInitialFocus = useCallback(() => {
     if (initialFocus === 'manual') {
       return;
     }
     if (initialFocus === 'auto') {
-      // An explicit `initialFocus` has not be set, so use automatic focus
+      // An explicit `initialFocus` has not been set, so use automatic focus
       // handling. Modern accessibility guidance is to focus the dialog itself
       // rather than trying to be smart about focusing a particular control
       // within the dialog.
@@ -109,8 +124,31 @@ const DialogNext = function Dialog({
     } else {
       // Fall back to focusing the modal itself
       modalRef.current?.focus();
-      return;
     }
+  }, [initialFocus, modalRef]);
+
+  useClickAway(modalRef, closeHandler, {
+    enabled: closeOnClickAway,
+  });
+
+  useKeyPress(['Escape'], closeHandler, {
+    enabled: closeOnEscape,
+  });
+
+  useFocusAway(modalRef, closeHandler, {
+    enabled: closeOnFocusAway,
+  });
+
+  const dialogDescriptionId = useUniqueId('dialog-description');
+  const Wrapper = useMemo(
+    () => TransitionComponent ?? Fragment,
+    [TransitionComponent]
+  );
+
+  useEffect(() => {
+    // Trigger initial open animation
+    setVisible(true);
+    setInitialFocus();
 
     // We only want to run this effect once when the dialog is mounted.
     // eslint-disable-next-line react-hooks/exhaustive-deps
@@ -152,31 +190,44 @@ const DialogNext = function Dialog({
   );
 
   return (
-    <div
-      data-component="Dialog"
-      tabIndex={-1}
-      // NB: Role can be overridden with an HTML attribute; this is purposeful
-      role="dialog"
-      {...htmlAttributes}
-      className={classnames(
-        // Column-flex layout to constrain content to max-height
-        'flex flex-col',
-        classes
-      )}
-      ref={downcastRef(modalRef)}
+    <Wrapper
+      visible={visible}
+      onTransitionEnd={() => {
+        if (!visible) {
+          onClose?.();
+        } else {
+          // Once transition on a visible Dialog has finished, we re-check the
+          // initial focus
+          setInitialFocus();
+        }
+      }}
     >
-      <Panel
-        buttons={buttons}
-        fullWidthHeader={true}
-        icon={icon}
-        onClose={onClose}
-        paddingSize={paddingSize}
-        title={title}
-        scrollable={scrollable}
+      <div
+        data-component="Dialog"
+        tabIndex={-1}
+        // NB: Role can be overridden with an HTML attribute; this is purposeful
+        role="dialog"
+        {...htmlAttributes}
+        className={classnames(
+          // Column-flex layout to constrain content to max-height
+          'flex flex-col',
+          classes
+        )}
+        ref={downcastRef(modalRef)}
       >
-        {children}
-      </Panel>
-    </div>
+        <Panel
+          buttons={buttons}
+          fullWidthHeader={true}
+          icon={icon}
+          onClose={closeHandler}
+          paddingSize={paddingSize}
+          title={title}
+          scrollable={scrollable}
+        >
+          {children}
+        </Panel>
+      </div>
+    </Wrapper>
   );
 };