diff --git a/src/cleanup.luau b/src/cleanup.luau index 0bfb6c1..803be03 100644 --- a/src/cleanup.luau +++ b/src/cleanup.luau @@ -1,19 +1,43 @@ if not game then script = require "test/relative-string" end +local typeof = game and typeof or require "test/mock".typeof :: never local throw = require(script.Parent.throw) local graph = require(script.Parent.graph) local get_scope = graph.get_scope local add_cleanup = graph.add_cleanup -local function cleanup(callback: () -> ()) +local function helper(obj: any) + return + if typeof(obj) == "RBXScriptConnection" then function() obj:Disconnect() end + elseif typeof(obj) == "Instance" then function() obj:Destroy() end + elseif obj.destroy then function() obj:destroy() end + elseif obj.disconnect then function() obj:disconnect() end + elseif obj.Destroy then function() obj:Destroy() end + elseif obj.Disconnect then function() obj:Disconnect() end + else throw("cannot cleanup given object") +end + +local function cleanup(value: unknown) local scope = get_scope() if not scope then throw "cannot cleanup in a non-reactive scope" end; assert(scope) - add_cleanup(scope, callback) + if type(value) == "function" then + add_cleanup(scope, value :: () -> ()) + else + add_cleanup(scope, helper(value)) + end end -return cleanup +type Destroyable = { destroy: (any) -> () } | { Destroy: (any) -> () } +type Disconnectable = { disconnect: (any) -> () } | { Disconnect: (any) -> () } + +return cleanup :: + ( (callback: () -> ()) -> () ) & + ( (instance: Destroyable) -> () ) & + ( (connection: Disconnectable) -> () ) & + ( (instance: Instance) -> () ) & + ( (connection: RBXScriptConnection) -> () ) diff --git a/src/untrack.luau b/src/untrack.luau index d9ca120..86cdb7b 100644 --- a/src/untrack.luau +++ b/src/untrack.luau @@ -1,6 +1,5 @@ if not game then script = require "test/relative-string" end -local throw = require(script.Parent.throw) local graph = require(script.Parent.graph) type Node = graph.Node local get_scope = graph.get_scope diff --git a/test/tests.luau b/test/tests.luau index 8a11df1..73ea0b7 100644 --- a/test/tests.luau +++ b/test/tests.luau @@ -691,6 +691,7 @@ TEST("effect()", wrap_root(function() end)) TEST("cleanup()", wrap_root(function() + local root = vide.root local source = vide.source local effect = vide.effect local cleanup = vide.cleanup @@ -749,6 +750,25 @@ TEST("cleanup()", wrap_root(function() src(3) CHECK(testkit.seq(queue, { 1, 2, 1, 2 })) end + + do CASE "cleanup objects" + local ran = {} + + root(function(destroy) + effect(function() + cleanup { disconnect = function() ran.disconnect = true end } + cleanup { Disconnect = function() ran.Disconnect = true end } + cleanup { destroy = function() ran.destroy = true end } + cleanup { Destroy = function() ran.Destroy = true end } + destroy() + end) + end) + + CHECK(ran.disconnect) + CHECK(ran.Disconnect) + CHECK(ran.destroy) + CHECK(ran.Destroy) + end end)) TEST("create()", wrap_root(function()