From 4fcc0598ceb7bdd9de182960f6abae2635aef01f Mon Sep 17 00:00:00 2001
From: Leroy Korterink <leroy@mediamonks.com>
Date: Mon, 9 Oct 2023 10:42:56 +0200
Subject: [PATCH] #182 Only log missing context warning in public API

---
 .../TransitionPresence.context.ts             | 13 +-------
 .../TransitionPresence/TransitionPresence.tsx |  4 +--
 .../react-transition-presence/src/index.ts    |  9 ++++--
 .../src/useBeforeUnmount/useBeforeUnmount.ts  | 30 +++++++++++++++++--
 4 files changed, 36 insertions(+), 20 deletions(-)

diff --git a/packages/react-transition-presence/src/TransitionPresence/TransitionPresence.context.ts b/packages/react-transition-presence/src/TransitionPresence/TransitionPresence.context.ts
index 98b92df..9d2164d 100644
--- a/packages/react-transition-presence/src/TransitionPresence/TransitionPresence.context.ts
+++ b/packages/react-transition-presence/src/TransitionPresence/TransitionPresence.context.ts
@@ -1,17 +1,6 @@
-import { createContext, useContext, type RefObject } from 'react';
+import { createContext, type RefObject } from 'react';
 import type { BeforeUnmountCallback } from '../useBeforeUnmount/useBeforeUnmount.js';
 
 export type TransitionPresenceContextType = Set<RefObject<BeforeUnmountCallback>> | undefined;
 
 export const TransitionPresenceContext = createContext<TransitionPresenceContextType>(undefined);
-
-export function useTransitionPresence(): TransitionPresenceContextType {
-  const context = useContext(TransitionPresenceContext);
-
-  if (context === undefined) {
-    // eslint-disable-next-line no-console
-    console.warn('Component is not rendered in the context of a TransitionPresence');
-  }
-
-  return context;
-}
diff --git a/packages/react-transition-presence/src/TransitionPresence/TransitionPresence.tsx b/packages/react-transition-presence/src/TransitionPresence/TransitionPresence.tsx
index 9cf4cf9..e4c666a 100644
--- a/packages/react-transition-presence/src/TransitionPresence/TransitionPresence.tsx
+++ b/packages/react-transition-presence/src/TransitionPresence/TransitionPresence.tsx
@@ -10,7 +10,7 @@ import {
 import { childrenAreEqual } from '../_utils/childrenAreEqual.js';
 import { tick } from '../_utils/tick.js';
 import {
-  useBeforeUnmount,
+  useTransitionPresenceBeforeUnmount,
   type BeforeUnmountCallback,
 } from '../useBeforeUnmount/useBeforeUnmount.js';
 import { TransitionPresenceContext } from './TransitionPresence.context.js';
@@ -104,7 +104,7 @@ export function TransitionPresence({
   ]);
 
   // Apply same effect when TransitionPresence in tree updates
-  useBeforeUnmount(beforeUnmountPreviousChildren);
+  useTransitionPresenceBeforeUnmount(beforeUnmountPreviousChildren);
 
   return (
     <TransitionPresenceContext.Provider value={beforeUnmountCallbacks}>
diff --git a/packages/react-transition-presence/src/index.ts b/packages/react-transition-presence/src/index.ts
index b4561fa..5dbd313 100644
--- a/packages/react-transition-presence/src/index.ts
+++ b/packages/react-transition-presence/src/index.ts
@@ -1,3 +1,6 @@
-export * from './TransitionPresence/TransitionPresence.js';
-export * from './TransitionPresence/TransitionPresence.context.js';
-export * from './useBeforeUnmount/useBeforeUnmount.js';
+export { TransitionPresenceContext } from './TransitionPresence/TransitionPresence.context.js';
+export { TransitionPresence } from './TransitionPresence/TransitionPresence.js';
+export {
+  type BeforeUnmountCallback,
+  useBeforeUnmount,
+} from './useBeforeUnmount/useBeforeUnmount.js';
diff --git a/packages/react-transition-presence/src/useBeforeUnmount/useBeforeUnmount.ts b/packages/react-transition-presence/src/useBeforeUnmount/useBeforeUnmount.ts
index 01bc9f1..faa7d8c 100644
--- a/packages/react-transition-presence/src/useBeforeUnmount/useBeforeUnmount.ts
+++ b/packages/react-transition-presence/src/useBeforeUnmount/useBeforeUnmount.ts
@@ -1,6 +1,6 @@
 import { useRefValue } from '@mediamonks/react-hooks';
-import { useEffect } from 'react';
-import { useTransitionPresence } from '../TransitionPresence/TransitionPresence.context.js';
+import { useContext, useEffect } from 'react';
+import { TransitionPresenceContext } from '../index.js';
 
 export type BeforeUnmountCallback = (
   abortSignal: AbortSignal,
@@ -12,7 +12,31 @@ PromiseLike<unknown> | void;
  * TransitionPresence boundary
  */
 export function useBeforeUnmount(callback: BeforeUnmountCallback): void {
-  const transitionPresence = useTransitionPresence();
+  const transitionPresence = useContext(TransitionPresenceContext);
+  const callbackRef = useRefValue(callback);
+
+  if (transitionPresence === undefined) {
+    // eslint-disable-next-line no-console
+    console.warn('Component is not rendered in the context of a TransitionPresence');
+  }
+
+  useEffect(() => {
+    queueMicrotask(() => {
+      transitionPresence?.add(callbackRef);
+    });
+
+    return () => {
+      transitionPresence?.delete(callbackRef);
+    };
+  }, [transitionPresence, callbackRef]);
+}
+
+/**
+ * useBeforeUnmount without the warning, this should only be used within the
+ * <TransitionPresence> component in this package.
+ */
+export function useTransitionPresenceBeforeUnmount(callback: BeforeUnmountCallback): void {
+  const transitionPresence = useContext(TransitionPresenceContext);
   const callbackRef = useRefValue(callback);
 
   useEffect(() => {