#182 Only log missing context warning in public API

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;
-}

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}>

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';

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(() => {