feat: support ctx.effect() returning object

src/scope.ts

@@ -7,7 +7,7 @@ declare module './context' {
   export interface Context {
     scope: EffectScope<this>
     runtime: MainScope<this>
-    effect(callback: () => () => void): () => void
+    effect<T extends DisposableLike>(callback: () => T): T
     /** @deprecated use `ctx.effect()` instead */
     collect(label: string, callback: () => void): () => void
     accept(callback?: (config: this['config']) => void | boolean, options?: AcceptOptions): () => boolean
@@ -18,6 +18,8 @@ declare module './context' {
 
 export type Disposable = () => void
 
+export type DisposableLike = Disposable | { dispose: Disposable }
+
 export interface AcceptOptions {
   passive?: boolean
   immediate?: boolean
@@ -86,15 +88,19 @@ export abstract class EffectScope<C extends Context = Context> {
     throw new CordisError('INACTIVE_EFFECT')
   }
 
-  effect(callback: () => () => void) {
+  effect(callback: () => DisposableLike) {
     this.assertActive()
-    const disposeRaw = callback()
-    const dispose = () => {
-      remove(this.disposables, dispose)
-      disposeRaw()
+    const result = callback()
+    const original = typeof result === 'function' ? result : result.dispose.bind(result)
+    const wrapped = () => {
+      // make sure the original callback is not called twice
+      if (!remove(this.disposables, wrapped)) return
+      return original()
     }
-    this.disposables.push(dispose)
-    return dispose
+    this.disposables.push(wrapped)
+    if (typeof result === 'function') return wrapped
+    result.dispose = wrapped
+    return result
   }
 
   collect(label: string, callback: () => any) {

tests/dispose.spec.ts

@@ -1,7 +1,7 @@
 import { Context } from '../src'
 import { expect } from 'chai'
 import { describe, mock, test } from 'node:test'
-import { noop } from 'cosmokit'
+import { noop, remove } from 'cosmokit'
 import { event, getHookSnapshot } from './utils'
 
 describe('Disposables', () => {
@@ -115,4 +115,40 @@ describe('Disposables', () => {
     expect(callback.mock.calls).to.have.length(2)
     expect(root.state.disposables.length).to.equal(length)
   })
+
+  test('ctx.effect()', async () => {
+    const root = new Context()
+    const dispose = mock.fn(noop)
+    const items: Item[] = []
+
+    class Item {
+      constructor() {
+        items.push(this)
+      }
+
+      dispose() {
+        dispose()
+        remove(items, this)
+      }
+    }
+
+    const item1 = root.effect(() => new Item())
+    const item2 = root.effect(() => new Item())
+    expect(item1).instanceof(Item)
+    expect(item2).instanceof(Item)
+    expect(dispose.mock.calls).to.have.length(0)
+    expect(items).to.have.length(2)
+
+    item1.dispose()
+    expect(dispose.mock.calls).to.have.length(1)
+    expect(items).to.have.length(1)
+
+    item1.dispose()
+    expect(dispose.mock.calls).to.have.length(1)
+    expect(items).to.have.length(1)
+
+    item2.dispose()
+    expect(dispose.mock.calls).to.have.length(2)
+    expect(items).to.have.length(0)
+  })
 })