@@ -40,31 +40,32 @@ def is_running(self) -> bool:
4040 running = self .task and not self .task .done ()
4141 return running
4242
43- def run_async (self , coro_func : Callable [[], asyncio . Future ] ):
43+ def run (self , func : Callable , * args , ** kwargs ):
4444 """
45- Run a pure async function in the background .
45+ Automatically detect function type and run appropriately .
4646
4747 Args:
48- coro_func: a non-blocking async function
48+ func: async or sync function
49+ *args, **kwargs: arguments to pass to func
4950 """
50- if not inspect .iscoroutinefunction (coro_func ):
51- raise TypeError ("run_async expects an async def function (not called yet)." )
52- self ._start (coro_func )
51+ if inspect .iscoroutinefunction (func ):
52+ self .run_async (func , * args , ** kwargs )
53+ else :
54+ self .run_blocking (func , * args , ** kwargs )
5355
54- def run_blocking (self , blocking_func : Callable [[], any ] ):
56+ def run_blocking (self , blocking_func : Callable , * args , ** kwargs ):
5557 """
5658 Run a blocking (sync) function via asyncio.to_thread.
5759
5860 Args:
5961 blocking_func: a regular function doing I/O or CPU-heavy work
62+ *args, **kwargs: arguments to pass to blocking_func
6063 """
61- if inspect .iscoroutinefunction (blocking_func ) or inspect .iscoroutine (
62- blocking_func
63- ):
64+ if inspect .iscoroutinefunction (blocking_func ) or inspect .iscoroutine (blocking_func ):
6465 raise TypeError ("run_blocking expects a sync function, not async." )
6566
6667 async def wrapper ():
67- await asyncio .to_thread (blocking_func )
68+ await asyncio .to_thread (blocking_func , * args , ** kwargs )
6869
6970 self ._start (wrapper )
7071
@@ -81,19 +82,29 @@ def run_async_wrapping_blocking(self, coro_func: Callable[[], asyncio.Future]):
8182 )
8283 self ._start (coro_func )
8384
84- def _start (self , coro_func : Callable [[], asyncio .Future ]):
85+ def run_async (self , coro_func : Callable [..., asyncio .Future ], * args , ** kwargs ):
86+ """
87+ Run a pure async function in the background.
88+
89+ Args:
90+ coro_func: a non-blocking async function
91+ *args, **kwargs: arguments to pass to coro_func
92+ """
93+ if not inspect .iscoroutinefunction (coro_func ):
94+ raise TypeError ("run_async expects an async def function (not called yet)." )
95+ self ._start (coro_func , * args , ** kwargs )
96+
97+ def _start (self , coro_func : Callable [..., asyncio .Future ], * args , ** kwargs ):
8598 self .cancel_running ()
8699
87100 async def wrapped ():
88101 try :
89102 if self .progress :
90103 self .progress .reset ()
91104 self .progress .set_description ("Working..." )
92- await asyncio .wait_for (coro_func (), timeout = self .timeout )
105+ await asyncio .wait_for (coro_func (* args , ** kwargs ), timeout = self .timeout )
93106 except asyncio .TimeoutError :
94- self .log .log (
95- "❌" , "timeout" , "Task exceeded timeout — possible blocking code?"
96- )
107+ self .log .log ("❌" , "timeout" , "Task exceeded timeout — possible blocking code?" )
97108 except asyncio .CancelledError :
98109 self .log .log ("⚠️" , "cancelled" , "Task was cancelled." )
99110 except Exception as ex :
0 commit comments