diff --git a/src/scope.ts b/src/scope.ts index 3ae47d9..dacde11 100644 --- a/src/scope.ts +++ b/src/scope.ts @@ -7,7 +7,7 @@ declare module './context' { export interface Context { scope: EffectScope runtime: MainScope - effect(callback: () => () => void): () => void + effect(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 { 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) { diff --git a/tests/dispose.spec.ts b/tests/dispose.spec.ts index ee79b8f..b6f4594 100644 --- a/tests/dispose.spec.ts +++ b/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) + }) })