Skip to content

Commit 671ae40

Browse files
committed
feat: allow custom exit function through dependency injection
1 parent 56a9a93 commit 671ae40

File tree

1 file changed

+11
-7
lines changed

1 file changed

+11
-7
lines changed

src/apify/_actor.py

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -231,17 +231,19 @@ async def exit(
231231
event_listeners_timeout: timedelta | None = EVENT_LISTENERS_TIMEOUT,
232232
status_message: str | None = None,
233233
cleanup_timeout: timedelta = timedelta(seconds=30),
234+
sys_exit: Callable = sys.exit,
234235
) -> None:
235236
"""Exit the Actor instance.
236237
237238
This stops the Actor instance. It cancels all the intervals for regularly sending `PERSIST_STATE` events,
238239
sends a final `PERSIST_STATE` event, waits for all the event listeners to finish, and stops the event manager.
239240
240241
Args:
241-
exit_code: The exit code with which the Actor should fail (defaults to `0`).
242+
exit_code: The exit code with which the program should end (defaults to `0`).
242243
event_listeners_timeout: How long should the Actor wait for Actor event listeners to finish before exiting.
243244
status_message: The final status message that the Actor should display.
244245
cleanup_timeout: How long we should wait for event listeners.
246+
sys_exit: A function to use to terminate the program (defaults to `sys.exit`)
245247
"""
246248
self._raise_if_not_initialized()
247249

@@ -267,29 +269,31 @@ async def finalize() -> None:
267269
self._is_initialized = False
268270

269271
if is_running_in_ipython():
270-
self.log.debug(f'Not calling sys.exit({exit_code}) because Actor is running in IPython')
272+
self.log.debug(f'Not calling sys_exit({exit_code}) because Actor is running in IPython')
271273
elif os.getenv('PYTEST_CURRENT_TEST', default=False): # noqa: PLW1508
272-
self.log.debug(f'Not calling sys.exit({exit_code}) because Actor is running in an unit test')
274+
self.log.debug(f'Not calling sys_exit({exit_code}) because Actor is running in an unit test')
273275
elif hasattr(asyncio, '_nest_patched'):
274-
self.log.debug(f'Not calling sys.exit({exit_code}) because Actor is running in a nested event loop')
276+
self.log.debug(f'Not calling sys_exit({exit_code}) because Actor is running in a nested event loop')
275277
else:
276-
sys.exit(exit_code)
278+
sys_exit(exit_code)
277279

278280
async def fail(
279281
self,
280282
*,
281283
exit_code: int = 1,
282284
exception: BaseException | None = None,
283285
status_message: str | None = None,
286+
sys_exit: Callable = sys.exit,
284287
) -> None:
285288
"""Fail the Actor instance.
286289
287290
This performs all the same steps as Actor.exit(), but it additionally sets the exit code to `1` (by default).
288291
289292
Args:
290-
exit_code: The exit code with which the Actor should fail (defaults to `1`).
293+
exit_code: The exit code with which the program should fail (defaults to `1`).
291294
exception: The exception with which the Actor failed.
292295
status_message: The final status message that the Actor should display.
296+
sys_exit: A function to use to terminate the program (defaults to `sys.exit`)
293297
"""
294298
self._raise_if_not_initialized()
295299

@@ -298,7 +302,7 @@ async def fail(
298302
if exception and not is_running_in_ipython():
299303
self.log.exception('Actor failed with an exception', exc_info=exception)
300304

301-
await self.exit(exit_code=exit_code, status_message=status_message)
305+
await self.exit(exit_code=exit_code, status_message=status_message, sys_exit=sys_exit)
302306

303307
def new_client(
304308
self,

0 commit comments

Comments
 (0)