Extend useRouter to also use transitions

example/app/page.js

@@ -1,17 +1,38 @@
-import { Link } from 'next-view-transitions'
+'use client'
+
+import { Link, useTransitionRouter } from 'next-view-transitions'
+import { useState } from "react";
 
 export default function Page() {
+  const [withCustomTransition, setWithCustomTransition] = useState(false)
+  const router = useTransitionRouter();
+
+  const routerNavigate = () => {
+    router.push('/demo', {
+        onTransitionReady: withCustomTransition ? slideInOut: undefined,
+    });
+  }
+
   return (
     <div>
       <h2>
         <span className='demo'>Demo</span>
       </h2>
-      <p>
-        <Link href='/demo'>Go to /demo →</Link>
-      </p>
-      <h2>Disclaimer</h2>
-      <p>
-        This library is aimed at basic use cases of View Transitions and Next.js
+        <p>
+            <Link href='/demo'>Go to /demo →</Link>
+        </p>
+        <p>
+            <a onClick={routerNavigate}>Go to /demo with router.push →</a>
+        </p>
+        <p>
+            <label>
+                <input type="checkbox" onChange={() => setWithCustomTransition((prev) => !prev)}/>{' '}
+                custom transition
+            </label>
+        </p>
+        <h2>Disclaimer</h2>
+        <p>
+            This library is aimed at basic use cases of View Transitions and Next.js
         App Router. With more complex applications and use cases like concurrent
         rendering, Suspense and streaming, new primitives and APIs still need to
         be developed into the core of React and Next.js in the future (
@@ -76,3 +97,44 @@ export default function Component() {
     </div>
   )
 }
+
+function slideInOut() {
+    document.documentElement.animate(
+        [
+            {
+                opacity: 1,
+                transform: 'translate(0, 0)',
+            },
+            {
+                opacity: 0,
+                transform: 'translate(-100%, 0)',
+            },
+        ],
+        {
+            duration: 500,
+            easing: 'ease-in-out',
+            fill: 'forwards',
+            pseudoElement: '::view-transition-old(root)',
+        }
+    );
+
+    document.documentElement.animate(
+        [
+            {
+                opacity: 0,
+                transform: 'translate(100%, 0)',
+            },
+            {
+                opacity: 1,
+                transform: 'translate(0, 0)',
+            },
+        ],
+        {
+            duration: 500,
+            easing: 'ease-in-out',
+            fill: 'forwards',
+            pseudoElement: '::view-transition-new(root)',
+        }
+    );
+}
+

src/index.ts

@@ -2,3 +2,4 @@
 
 export { Link } from './link'
 export { ViewTransitions } from './transition-context'
+export { useTransitionRouter } from './use-transition-router'

src/link.tsx

@@ -1,7 +1,6 @@
 import NextLink from 'next/link'
-import { useRouter } from 'next/navigation'
-import { startTransition, useCallback } from 'react'
-import { useSetFinishViewTransition } from './transition-context'
+import { useTransitionRouter } from './use-transition-router'
+import { useCallback } from 'react'
 
 // copied from https://github.com/vercel/next.js/blob/66f8ffaa7a834f6591a12517618dce1fd69784f6/packages/next/src/client/link.tsx#L180-L191
 function isModifiedEvent(event: React.MouseEvent): boolean {
@@ -38,8 +37,7 @@ function shouldPreserveDefault(
 // to navigate, and trigger a view transition.
 
 export function Link(props: React.ComponentProps<typeof NextLink>) {
-  const router = useRouter()
-  const finishViewTransition = useSetFinishViewTransition()
+  const router = useTransitionRouter()
 
   const { href, as, replace, scroll } = props
   const onClick = useCallback(
@@ -55,19 +53,8 @@ export function Link(props: React.ComponentProps<typeof NextLink>) {
 
         e.preventDefault()
 
-        // @ts-ignore
-        document.startViewTransition(
-          () =>
-            new Promise<void>((resolve) => {
-              startTransition(() => {
-                // copied from https://github.com/vercel/next.js/blob/66f8ffaa7a834f6591a12517618dce1fd69784f6/packages/next/src/client/link.tsx#L231-L233
-                router[replace ? 'replace' : 'push'](as || href, {
-                  scroll: scroll ?? true,
-                })
-                finishViewTransition(() => resolve)
-              })
-            })
-        )
+        const navigate = replace ? router.replace : router.push
+        navigate(as || href, { scroll: scroll ?? true })
       }
     },
     [props.onClick, href, as, replace, scroll]

src/use-transition-router.ts

@@ -0,0 +1,66 @@
+import { useRouter as useNextRouter } from 'next/navigation'
+import {startTransition, useCallback, useMemo} from "react";
+import { useSetFinishViewTransition } from "./transition-context";
+import {
+  AppRouterInstance,
+  NavigateOptions
+} from "next/dist/shared/lib/app-router-context.shared-runtime";
+
+export type TransitionOptions = {
+  onTransitionReady?: () => void;
+};
+
+type NavigateOptionsWithTransition = NavigateOptions & TransitionOptions;
+
+export type TransitionRouter = AppRouterInstance & {
+  push: (href: string, options?: NavigateOptionsWithTransition) => void;
+  replace: (href: string, options?: NavigateOptionsWithTransition) => void;
+};
+
+export function useTransitionRouter() {
+  const router = useNextRouter()
+  const finishViewTransition = useSetFinishViewTransition()
+
+  const triggerTransition = useCallback((cb: () => void, { onTransitionReady }: TransitionOptions = {}) => {
+    // @ts-ignore
+    const transition = document.startViewTransition(
+      () =>
+        new Promise<void>((resolve) => {
+          startTransition(() => {
+            cb();
+            finishViewTransition(() => resolve)
+          })
+        })
+    )
+
+    if (onTransitionReady) {
+      transition.ready.then(onTransitionReady);
+    }
+  }, [])
+
+  const push = useCallback((
+    href: string,
+    { onTransitionReady, ...options }: NavigateOptionsWithTransition = {}
+  ) => {
+    triggerTransition(() => router.push(href, options), {
+      onTransitionReady
+    })
+  }, [triggerTransition, router])
+
+  const replace = useCallback((
+    href: string,
+    { onTransitionReady, ...options }: NavigateOptionsWithTransition = {}
+  ) => {
+    triggerTransition(() => router.replace(href, options), {
+      onTransitionReady
+    });
+  }, [triggerTransition, router]);
+
+  return useMemo<TransitionRouter>(
+    () => ({
+      ...router,
+      push,
+      replace,
+    }),
+    [push, replace, router]);
+}