@@ -137,21 +137,21 @@ def use_effect(
137
137
hook = current_hook ()
138
138
dependencies = _try_to_infer_closure_values (function , dependencies )
139
139
memoize = use_memo (dependencies = dependencies )
140
- unmount_func : Ref [_EffectCleanFunc | None ] = use_ref (None )
140
+ cleanup_func : Ref [_EffectCleanFunc | None ] = use_ref (None )
141
141
142
142
def decorator (func : _SyncEffectFunc ) -> None :
143
143
async def effect (stop : asyncio .Event ) -> None :
144
- if unmount_func .current :
145
- unmount_func .current ()
146
- unmount_func .current = None
144
+ if cleanup_func .current :
145
+ cleanup_func .current ()
146
+ cleanup_func .current = None
147
147
148
148
# Execute the effect and store the clean-up function
149
- unmount = unmount_func .current = func ()
149
+ cleanup_func .current = func ()
150
150
151
151
# Run the clean-up function when the effect is stopped
152
152
await stop .wait ()
153
- if unmount :
154
- unmount ()
153
+ if cleanup_func . current :
154
+ cleanup_func . current ()
155
155
156
156
return memoize (lambda : hook .add_effect (effect ))
157
157
@@ -179,54 +179,34 @@ def use_async_effect(
179
179
def use_async_effect (
180
180
function : _AsyncEffectFunc | None = None ,
181
181
dependencies : Sequence [Any ] | ellipsis | None = ...,
182
+ shutdown_timeout : float = 0.1 ,
182
183
) -> Callable [[_AsyncEffectFunc ], None ] | None :
183
- """See the full :ref:`Use Effect` docs for details
184
-
185
- Parameters:
186
- function:
187
- Applies the effect and can return a clean-up function
188
- dependencies:
189
- Dependencies for the effect. The effect will only trigger if the identity
190
- of any value in the given sequence changes (i.e. their :func:`id` is
191
- different). By default these are inferred based on local variables that are
192
- referenced by the given function.
193
-
194
- Returns:
195
- If not function is provided, a decorator. Otherwise ``None``.
196
- """
197
184
hook = current_hook ()
198
185
dependencies = _try_to_infer_closure_values (function , dependencies )
199
186
memoize = use_memo (dependencies = dependencies )
200
- unmount_func : Ref [_EffectCleanFunc | None ] = use_ref (None )
187
+ cleanup_func : Ref [_EffectCleanFunc | None ] = use_ref (None )
201
188
202
189
def decorator (func : _AsyncEffectFunc ) -> None :
203
- def sync_executor () -> _EffectCleanFunc | None :
204
- task = asyncio .create_task (func ())
205
-
206
- def unmount_executor () -> None :
207
- if not task .cancel ():
208
- try :
209
- unmount = task .result ()
210
- except asyncio .CancelledError :
211
- pass
212
- else :
213
- if unmount :
214
- unmount ()
215
-
216
- return unmount_executor
217
-
218
190
async def effect (stop : asyncio .Event ) -> None :
219
- if unmount_func .current :
220
- unmount_func .current ()
221
- unmount_func .current = None
191
+ if cleanup_func .current :
192
+ cleanup_func .current ()
193
+ cleanup_func .current = None
222
194
223
- # Execute the effect and store the clean-up function
224
- unmount = unmount_func . current = sync_executor ( )
195
+ # Execute the effect in a background task
196
+ task = asyncio . create_task ( func () )
225
197
226
- # Run the clean-up function when the effect is stopped
198
+ # Wait until the effect is stopped
227
199
await stop .wait ()
228
- if unmount :
229
- unmount ()
200
+
201
+ # Try to fetch the results of the task
202
+ results , _ = await asyncio .wait ([task ], timeout = shutdown_timeout )
203
+ if results :
204
+ cleanup_func .current = results .pop ().result ()
205
+ if cleanup_func .current :
206
+ cleanup_func .current ()
207
+
208
+ # Cancel the task if it's still running
209
+ task .cancel ()
230
210
231
211
return memoize (lambda : hook .add_effect (effect ))
232
212
0 commit comments