From d4420bb9734aa3e1539d3937b1e027e9f51d65a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=89=E5=92=B2=E6=99=BA=E5=AD=90=20Kevin=20Deng?=
 <sxzz@sxzz.moe>
Date: Wed, 19 Mar 2025 04:19:27 +0800
Subject: [PATCH] feat: add `onYield` hook

---
 src/index.ts       | 43 ++++++++++++++++++++++++++++++-------------
 src/types.ts       |  6 +++++-
 test/index.test.ts | 28 +++++++++++++++++++++++++++-
 3 files changed, 62 insertions(+), 15 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index e7a9557..67490e3 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,4 +1,4 @@
-import type { QuansyncAwaitableGenerator, QuansyncFn, QuansyncGenerator, QuansyncGeneratorFn, QuansyncInput, QuansyncInputObject } from './types'
+import type { QuansyncAwaitableGenerator, QuansyncFn, QuansyncGenerator, QuansyncGeneratorFn, QuansyncInput, QuansyncInputObject, QuansyncOptions } from './types'
 
 export * from './types'
 
@@ -59,11 +59,16 @@ function unwrapYield(value: any, isAsync?: boolean): any {
   return value
 }
 
-function iterateSync<Return>(generator: QuansyncGenerator<Return, unknown>): Return {
+const DEFAULT_ON_YIELD = (value: any): any => value
+
+function iterateSync<Return>(
+  generator: QuansyncGenerator<Return, unknown>,
+  onYield: QuansyncOptions['onYield'] = DEFAULT_ON_YIELD,
+): Return {
   let current = generator.next()
   while (!current.done) {
     try {
-      current = generator.next(unwrapYield(current.value))
+      current = generator.next(unwrapYield(onYield(current.value, false)))
     }
     catch (err) {
       current = generator.throw(err)
@@ -72,11 +77,14 @@ function iterateSync<Return>(generator: QuansyncGenerator<Return, unknown>): Ret
   return unwrapYield(current.value)
 }
 
-async function iterateAsync<Return>(generator: QuansyncGenerator<Return, unknown>): Promise<Return> {
+async function iterateAsync<Return>(
+  generator: QuansyncGenerator<Return, unknown>,
+  onYield: QuansyncOptions['onYield'] = DEFAULT_ON_YIELD,
+): Promise<Return> {
   let current = generator.next()
   while (!current.done) {
     try {
-      current = generator.next(await unwrapYield(current.value, true))
+      current = generator.next(await unwrapYield(onYield(current.value, true), true))
     }
     catch (err) {
       current = generator.throw(err)
@@ -87,14 +95,15 @@ async function iterateAsync<Return>(generator: QuansyncGenerator<Return, unknown
 
 function fromGeneratorFn<Return, Args extends any[]>(
   generatorFn: QuansyncGeneratorFn<Return, Args>,
+  options?: QuansyncOptions,
 ): QuansyncFn<Return, Args> {
   return fromObject({
     name: generatorFn.name,
     async(...args) {
-      return iterateAsync(generatorFn.apply(this, args))
+      return iterateAsync(generatorFn.apply(this, args), options?.onYield)
     },
     sync(...args) {
-      return iterateSync(generatorFn.apply(this, args))
+      return iterateSync(generatorFn.apply(this, args), options?.onYield)
     },
   })
 }
@@ -103,14 +112,22 @@ function fromGeneratorFn<Return, Args extends any[]>(
  * Creates a new Quansync function, a "superposition" between async and sync.
  */
 export function quansync<Return, Args extends any[] = []>(
-  options: QuansyncInput<Return, Args> | Promise<Return>,
+  input: QuansyncInputObject<Return, Args>,
+): QuansyncFn<Return, Args>
+export function quansync<Return, Args extends any[] = []>(
+  input: QuansyncGeneratorFn<Return, Args> | Promise<Return>,
+  options?: QuansyncOptions,
+): QuansyncFn<Return, Args>
+export function quansync<Return, Args extends any[] = []>(
+  input: QuansyncInput<Return, Args> | Promise<Return>,
+  options?: QuansyncOptions,
 ): QuansyncFn<Return, Args> {
-  if (isThenable(options))
-    return fromPromise<Return>(options)
-  if (typeof options === 'function')
-    return fromGeneratorFn(options)
+  if (isThenable(input))
+    return fromPromise<Return>(input)
+  if (typeof input === 'function')
+    return fromGeneratorFn(input, options)
   else
-    return fromObject(options)
+    return fromObject(input)
 }
 
 /**
diff --git a/src/types.ts b/src/types.ts
index c089de7..b868009 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -1,4 +1,8 @@
-export interface QuansyncInputObject<Return, Args extends any[]> {
+export interface QuansyncOptions {
+  onYield?: (value: any, isAsync: boolean) => any
+}
+
+export interface QuansyncInputObject<Return, Args extends any[]> extends QuansyncOptions {
   name?: string
   sync: (...args: Args) => Return
   async: (...args: Args) => Promise<Return>
diff --git a/test/index.test.ts b/test/index.test.ts
index 5c27443..b570ffb 100644
--- a/test/index.test.ts
+++ b/test/index.test.ts
@@ -1,4 +1,4 @@
-import { expect, it } from 'vitest'
+import { expect, it, vi } from 'vitest'
 import { quansync, toGenerator } from '../src'
 import { quansync as quansyncMacro } from '../src/macro'
 
@@ -256,3 +256,29 @@ it('bind this', async () => {
   await expect(cls.async()).resolves.instanceOf(Cls)
   expect(cls.sync()).instanceOf(Cls)
 })
+
+it('call onYield hook', async () => {
+  const onYield = vi.fn(v => v + 1)
+  const run = quansync(function* () {
+    expect(yield 1).toBe(2)
+    expect(yield 2).toBe(3)
+  }, { onYield })
+
+  await run.async()
+  expect(onYield).toBeCalledTimes(2)
+  run.sync()
+  expect(onYield).toBeCalledTimes(4)
+
+  // custom error on promise
+  const run2 = quansync(function* () {
+    return yield Promise.resolve('foo')
+  }, {
+    onYield: (v, isAsync) => {
+      if (!isAsync && v instanceof Promise)
+        throw new TypeError('custom error')
+      return v
+    },
+  })
+  expect(() => run2.sync()).toThrowErrorMatchingInlineSnapshot(`[TypeError: custom error]`)
+  await expect(run2.async()).resolves.toMatchInlineSnapshot(`"foo"`)
+})