allow passing popen_kwargs

Author: Nanguage

executor/engine/__init__.py

@@ -1,7 +1,7 @@
 from .core import Engine, EngineSetting
 from .job import LocalJob, ThreadJob, ProcessJob
 
-__version__ = '0.2.3'
+__version__ = '0.2.4'
 
 __all__ = [
     'Engine', 'EngineSetting',

executor/engine/job/extend/subprocess.py

@@ -1,5 +1,4 @@
 import copy
-import shlex
 import typing as T
 import subprocess as subp
 from pathlib import Path
@@ -23,6 +22,7 @@ def SubprocessJob(
     wait_time_delta: float = 0.01,
     redirect_out_err: bool = False,
     target_dir: str = "$current_dir",
+    popen_kwargs: T.Optional[T.Dict[str, T.Any]] = None,
     **attrs
 ):
     """Create a job that runs a subprocess.
@@ -43,6 +43,7 @@ def SubprocessJob(
             Use '$cache_dir' to represent the cache directory of the job.
             Use '$current_dir' to represent the current directory of the job.
             Default is '$current_dir'.
+        popen_kwargs: The keyword arguments for subprocess.Popen.
         **attrs: Other attributes of the job.
     """
     class _SubprocessJob(base_class):  # type: ignore
@@ -107,21 +108,29 @@ def record_command():
                 with open(path_sh, 'w') as f:
                     f.write(cmd + "\n")
 
+            pkwargs = popen_kwargs or {}
+            pkwargs['cwd'] = target_dir
+
             if self.redirect_out_err:
                 path_stdout = cache_dir / 'stdout.txt'
                 path_stderr = cache_dir / 'stderr.txt'
 
                 def run_cmd():  # pragma: no cover
                     runner = ProcessRunner(cmd)
-                    runner.run(cwd=target_dir)
+                    runner.run(**pkwargs)
                     with open(path_stdout, 'w') as fo, \
                          open(path_stderr, 'w') as fe:
                         retcode = runner.write_stream_until_stop(fo, fe)
                     return retcode
             else:
                 def run_cmd():
-                    p = subp.Popen(shlex.split(cmd), cwd=target_dir)
-                    retcode = p.wait()
+                    runner = ProcessRunner(cmd)
+                    runner.run(
+                        capture_stdout=False,
+                        capture_stderr=False,
+                        **pkwargs
+                    )
+                    retcode = runner.proc.wait()
                     return retcode
 
             def func():

requirements.txt

@@ -1,4 +1,4 @@
-cmd2func>=0.1.6
+cmd2func>=0.1.9
 funcdesc>=0.1.2
 loky
 diskcache

tests/test_subp_job.py

@@ -134,3 +134,23 @@ async def submit_job():
         assert job2.result() == 0
 
     asyncio.run(submit_job())
+
+
+def test_passing_env():
+    engine = Engine()
+
+    async def submit_job():
+        import os
+        env = os.environ.copy()
+        env["A"] = "2"
+        job = SubprocessJob(
+            "python -c 'import os; print(os.getenv(\"A\"))'",
+            popen_kwargs={"env": env},
+            redirect_out_err=True
+        )
+        await engine.submit_async(job)
+        await job.join()
+        with open(job.cache_dir / 'stdout.txt') as f:
+            assert f.read().strip() == '2'
+
+    asyncio.run(submit_job())