From f292d32d4cc9a991259b7dc905a7c30c644611bc Mon Sep 17 00:00:00 2001
From: Soybean <soybeanjs@outlook.com>
Date: Wed, 13 May 2026 14:45:20 +0800
Subject: [PATCH] refactor(hooks): simplify useContext implementation and
 improve type definitions

---
 packages/hooks/src/use-context.ts | 99 ++++++++++---------------------
 1 file changed, 30 insertions(+), 69 deletions(-)

diff --git a/packages/hooks/src/use-context.ts b/packages/hooks/src/use-context.ts
index cea9164f..08a6a519 100644
--- a/packages/hooks/src/use-context.ts
+++ b/packages/hooks/src/use-context.ts
@@ -1,66 +1,31 @@
 import { inject, provide } from 'vue';
 
+type ContextName = string | { name: string; key: string | symbol };
+
+type ContextValue<T> = T extends (...args: any[]) => any ? ReturnType<T> : T;
+
+type ContextProvider<T> = T extends (...args: any[]) => any ? T : (arg: T) => T;
+
+type ContextConsumer<Context> = <N extends string | null | undefined = undefined>(
+  consumerName?: N,
+  defaultValue?: Context
+) => N extends null | undefined ? Context | null : Context;
+
 /**
- * Use context
+ * Creates a context provider and consumer pair.
  *
- * @example
- *   ```ts
- *   // there are three vue files: A.vue, B.vue, C.vue, and A.vue is the parent component of B.vue and C.vue
- *
- *   // context.ts
- *   import { ref } from 'vue';
- *   import { useContext } from '@sa/hooks';
- *
- *   export const [provideDemoContext, useDemoContext] = useContext('demo', () => {
- *     const count = ref(0);
- *
- *     function increment() {
- *       count.value++;
- *     }
- *
- *     function decrement() {
- *       count.value--;
- *     }
- *
- *     return {
- *       count,
- *       increment,
- *       decrement
- *     };
- *   })
- *   ``` // A.vue
- *   ```vue
- *   <template>
- *     <div>A</div>
- *   </template>
- *   <script setup lang="ts">
- *   import { provideDemoContext } from './context';
- *
- *   provideDemoContext();
- *   // const { increment } = provideDemoContext(); // also can control the store in the parent component
- *   </script>
- *   ``` // B.vue
- *   ```vue
- *   <template>
- *    <div>B</div>
- *   </template>
- *   <script setup lang="ts">
- *   import { useDemoContext } from './context';
- *
- *   const { count, increment } = useDemoContext();
- *   </script>
- *   ```;
- *
- *   // C.vue is same as B.vue
- *
- * @param contextName Context name
- * @param fn Context function
+ * @param contextName - The name of the context. This can be a string or an object with a `name` and `key` property.
+ * @param composable - An optional composable function that returns the context value. If not provided, the context value will be the first argument passed to the provider.
  */
-export default function useContext<Arguments extends Array<any>, T>(
-  contextName: string,
-  composable: (...args: Arguments) => T
+export default function useContext<T>(
+  contextName: ContextName,
+  composable?: T extends (...args: any[]) => any ? T : never
 ) {
-  const key = Symbol(contextName);
+  type Context = ContextValue<T>;
+
+  const name = typeof contextName === 'string' ? contextName : contextName.name;
+
+  const key = typeof contextName === 'string' ? Symbol(contextName) : contextName.key;
 
   /**
    * Injects the context value.
@@ -70,27 +35,23 @@ export default function useContext<Arguments extends Array<any>, T>(
    * @param defaultValue - The default value to return if the context is not provided.
    * @returns The context value.
    */
-  const useInject = <N extends string | null | undefined = undefined>(
-    consumerName?: N,
-    defaultValue?: T
-  ): N extends null | undefined ? T | null : T => {
-    const value = inject(key, defaultValue);
+  const useInject = (consumerName?: string | null, defaultValue?: any) => {
+    const value = inject(key, defaultValue) ?? null;
 
-    if (consumerName && !value) {
-      throw new Error(`\`${consumerName}\` must be used within \`${contextName}\``);
+    if (consumerName != null && value === null) {
+      throw new Error(`\`${consumerName}\` must be used within \`${name}\``);
     }
 
-    // @ts-expect-error - we want to return null if the value is undefined or null
-    return value || null;
+    return value;
   };
 
-  const useProvide = (...args: Arguments) => {
-    const value = composable(...args);
+  const useProvide = (...args: any[]) => {
+    const value = composable?.(...args) ?? args[0];
 
     provide(key, value);
 
     return value;
   };
 
-  return [useProvide, useInject] as const;
+  return [useProvide, useInject] as [ContextProvider<T>, ContextConsumer<Context>];
 }