From 990b1796a52b9ce1f12bb5225e7385bce68c3770 Mon Sep 17 00:00:00 2001
From: Robin Malfait <malfait.robin@gmail.com>
Date: Mon, 1 Jul 2024 15:24:01 +0200
Subject: [PATCH] Ensure `unmount` on `Dialog` works in combination with the
 `transition` prop on `DialogBackdrop` and `DialogPanel` components (#3352)

* inherit `unmount` from `Dialog` in `DialogBackdrop` and `DialogPanel` components

Only the `Dialog` accepts an `unmount` prop because it's the `Dialog`
that is conditionally rendered and the `DialogBackdrop` and
`DialogPanel` will conditionally show together with the `Dialog`.

However, now that the `Dialog` is wrapped in a `Transition` (which can
be unmounted) and the `DialogBackdrop` and `DialogPanel` will also be
wrapped in a `TransitionChild` (when the `transition` prop is passed)
then we do have to deal with the `unmount` state on the `TransitionChild`.

This is important because if you make the `Dialog` `unmount={false}`, then the
`DialogPanel` will still unmount because the `TransitionChild` is unmounting its
children. This now means that you will lose data (such as form state of inputs).

This commit solves that by inheriting the `unmount` state of the
`Dialog` in the `TransitionChild` wrappers such that they behave the way
you expect them to behave.

* update changelog
---
 packages/@headlessui-react/CHANGELOG.md         |  1 +
 .../src/components/dialog/dialog.tsx            | 17 +++++++++++------
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/packages/@headlessui-react/CHANGELOG.md b/packages/@headlessui-react/CHANGELOG.md
index 830d7f4036..5a86bafe99 100644
--- a/packages/@headlessui-react/CHANGELOG.md
+++ b/packages/@headlessui-react/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Fix prematurely added anchoring styles on `ListboxOptions` ([#3337](https://github.com/tailwindlabs/headlessui/pull/3337))
+- Ensure `unmount` on `Dialog` works in combination with the `transition` prop on `DialogBackdrop` and `DialogPanel` components ([#3352](https://github.com/tailwindlabs/headlessui/pull/3352))
 
 ## [2.1.1] - 2024-06-26
 
diff --git a/packages/@headlessui-react/src/components/dialog/dialog.tsx b/packages/@headlessui-react/src/components/dialog/dialog.tsx
index 1a9b7fe6d8..49fd061f4d 100644
--- a/packages/@headlessui-react/src/components/dialog/dialog.tsx
+++ b/packages/@headlessui-react/src/components/dialog/dialog.tsx
@@ -84,6 +84,7 @@ let DialogContext = createContext<
   | [
       {
         dialogState: DialogStates
+        unmount: boolean
         close(): void
         setTitleId(id: string | null): void
       },
@@ -121,6 +122,7 @@ let InternalDialog = forwardRefWithAs(function InternalDialog<
     role = 'dialog',
     autoFocus = true,
     __demoMode = false,
+    unmount = false,
     ...theirProps
   } = props
 
@@ -248,8 +250,8 @@ let InternalDialog = forwardRefWithAs(function InternalDialog<
   let [describedby, DescriptionProvider] = useDescriptions()
 
   let contextBag = useMemo<ContextType<typeof DialogContext>>(
-    () => [{ dialogState, close, setTitleId }, state],
-    [dialogState, state, close, setTitleId]
+    () => [{ dialogState, close, setTitleId, unmount }, state],
+    [dialogState, state, close, setTitleId, unmount]
   )
 
   let slot = useMemo(
@@ -265,6 +267,7 @@ let InternalDialog = forwardRefWithAs(function InternalDialog<
     'aria-modal': __demoMode ? undefined : dialogState === DialogStates.Open ? true : undefined,
     'aria-labelledby': state.titleId,
     'aria-describedby': describedby,
+    unmount,
   }
 
   let shouldMoveFocusInside = !useIsTouchDevice()
@@ -421,7 +424,7 @@ function PanelFn<TTag extends ElementType = typeof DEFAULT_PANEL_TAG>(
 ) {
   let internalId = useId()
   let { id = `headlessui-dialog-panel-${internalId}`, transition = false, ...theirProps } = props
-  let [{ dialogState }, state] = useDialogContext('Dialog.Panel')
+  let [{ dialogState, unmount }, state] = useDialogContext('Dialog.Panel')
   let panelRef = useSyncRefs(ref, state.panelRef)
 
   let slot = useMemo(
@@ -442,9 +445,10 @@ function PanelFn<TTag extends ElementType = typeof DEFAULT_PANEL_TAG>(
   }
 
   let Wrapper = transition ? TransitionChild : Fragment
+  let wrapperProps = transition ? { unmount } : {}
 
   return (
-    <Wrapper>
+    <Wrapper {...wrapperProps}>
       {render({
         ourProps,
         theirProps,
@@ -475,7 +479,7 @@ function BackdropFn<TTag extends ElementType = typeof DEFAULT_BACKDROP_TAG>(
   ref: Ref<HTMLElement>
 ) {
   let { transition = false, ...theirProps } = props
-  let [{ dialogState }] = useDialogContext('Dialog.Backdrop')
+  let [{ dialogState, unmount }] = useDialogContext('Dialog.Backdrop')
 
   let slot = useMemo(
     () => ({ open: dialogState === DialogStates.Open }) satisfies BackdropRenderPropArg,
@@ -485,9 +489,10 @@ function BackdropFn<TTag extends ElementType = typeof DEFAULT_BACKDROP_TAG>(
   let ourProps = { ref, 'aria-hidden': true }
 
   let Wrapper = transition ? TransitionChild : Fragment
+  let wrapperProps = transition ? { unmount } : {}
 
   return (
-    <Wrapper>
+    <Wrapper {...wrapperProps}>
       {render({
         ourProps,
         theirProps,