From 6e5f9b9cba530060b2dfaf844ef065dbe5873ec2 Mon Sep 17 00:00:00 2001
From: phi <phi.friday@gmail.com>
Date: Fri, 13 Sep 2024 13:07:55 +0900
Subject: [PATCH] fix: initfile args

---
 src/timeout_executor/executor.py | 4 +++-
 src/timeout_executor/result.py   | 6 +++++-
 src/timeout_executor/types.py    | 2 ++
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/timeout_executor/executor.py b/src/timeout_executor/executor.py
index afd57f0..9942a66 100644
--- a/src/timeout_executor/executor.py
+++ b/src/timeout_executor/executor.py
@@ -140,6 +140,7 @@ def _create_executor_args(
         self,
         input_file: Path | anyio.Path,
         output_file: Path | anyio.Path,
+        init_file: Path | anyio.Path | None,
         terminator: Terminator[P, T],
     ) -> ExecutorArgs[P, T]:
         """create executor args"""
@@ -149,6 +150,7 @@ def _create_executor_args(
             terminator=terminator,
             input_file=input_file,
             output_file=output_file,
+            init_file=init_file,
             timeout=self._timeout,
         )
 
@@ -177,7 +179,7 @@ def _init_process(
         """
         logger.debug("%r before init process", self, stacklevel=stacklevel)
         executor_args_builder = partial(
-            self._create_executor_args, input_file, output_file
+            self._create_executor_args, input_file, output_file, init_file
         )
         terminator = Terminator(executor_args_builder, self.callbacks)
         process = self._create_process(input_file, init_file, stacklevel=stacklevel + 1)
diff --git a/src/timeout_executor/result.py b/src/timeout_executor/result.py
index 3e893d9..aa5b42e 100644
--- a/src/timeout_executor/result.py
+++ b/src/timeout_executor/result.py
@@ -43,12 +43,14 @@ def __init__(
         self._executor_args = executor_args
         self._result = SENTINEL
 
-        input_file, output_file = (
+        input_file, output_file, init_file = (
             self._executor_args.input_file,
             self._executor_args.output_file,
+            self._executor_args.init_file,
         )
         self._input = anyio.Path(input_file)
         self._output = anyio.Path(output_file)
+        self._init = None if init_file is None else anyio.Path(init_file)
 
     @property
     def _func_name(self) -> str:
@@ -119,6 +121,8 @@ async def _load_output(self) -> T:
         async with anyio.create_task_group() as task_group:
             task_group.start_soon(self._output.unlink, True)  # noqa: FBT003
             task_group.start_soon(self._input.unlink, True)  # noqa: FBT003
+            if self._init is not None:
+                task_group.start_soon(self._init.unlink, True)  # noqa: FBT003
         await self._output.parent.rmdir()
         return await self._load_output()
 
diff --git a/src/timeout_executor/types.py b/src/timeout_executor/types.py
index f1f24f2..625808a 100644
--- a/src/timeout_executor/types.py
+++ b/src/timeout_executor/types.py
@@ -55,6 +55,8 @@ class ExecutorArgs(Generic[P, T]):
     """function args input file"""
     output_file: Path | anyio.Path
     """function result output file"""
+    init_file: Path | anyio.Path | None
+    """initializer file"""
     timeout: float
     """timeout"""