allow redirect stdout/err to specific file

Author: Nanguage

executor/engine/__init__.py

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

executor/engine/job/base.py

@@ -83,7 +83,7 @@ def __init__(
             name: T.Optional[str] = None,
             condition: T.Optional[Condition] = None,
             wait_time_delta: float = 0.01,
-            redirect_out_err: bool = False,
+            redirect_out_err: bool | str = False,
             change_dir: bool = False,
             **attrs
             ) -> None:
@@ -103,7 +103,8 @@ def __init__(
             wait_time_delta: The time delta between each
                 check of the condition.
             redirect_out_err: Whether to redirect the stdout
-                and stderr to the log.
+                and stderr to the log. If is a string, it will be
+                used as the path of the log file.
             change_dir: Whether to change the working directory
                 to the log directory.
             **attrs: The attributes of the job.
@@ -255,9 +256,13 @@ def process_func(self):
         change the dir, redirect the stdout and stderr
         before the actual run."""
         cache_dir = self.cache_dir.resolve()
-        if self.redirect_out_err and (not isinstance(self.func, CaptureOut)):
-            path_stdout = cache_dir / 'stdout.txt'
-            path_stderr = cache_dir / 'stderr.txt'
+        if (self.redirect_out_err is not False) and (not isinstance(self.func, CaptureOut)):
+            if isinstance(self.redirect_out_err, str):
+                path_stdout = Path(self.redirect_out_err)
+                path_stderr = Path(self.redirect_out_err)
+            else:
+                path_stdout = cache_dir / 'stdout.txt'
+                path_stderr = cache_dir / 'stderr.txt'
             self.func = CaptureOut(self.func, path_stdout, path_stderr)
         if self.change_dir:
             self.func = ChDir(self.func, cache_dir)

executor/engine/job/extend/subprocess.py

@@ -112,16 +112,30 @@ def record_command():
             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'
+            if self.redirect_out_err is not False:
+                if isinstance(self.redirect_out_err, str):
+                    path_stdout = Path(self.redirect_out_err)
+                    if not path_stdout.parent.exists():
+                        path_stdout.parent.mkdir(parents=True, exist_ok=True)
+                    path_stdout = self.redirect_out_err
+                    path_stderr = self.redirect_out_err
+                else:
+                    path_stdout = cache_dir / 'stdout.txt'
+                    path_stderr = cache_dir / 'stderr.txt'
 
                 def _run_cmd(runner: ProcessRunner):  # pragma: no cover
                     runner.run(**pkwargs)
-                    with open(path_stdout, 'w') as fo, \
-                         open(path_stderr, 'w') as fe:
-                        retcode = runner.write_stream_until_stop(
-                            fo, fe, flush_streams_each_time=True)
+                    if path_stdout == path_stderr:
+                        fo = open(path_stdout, 'a')
+                        fe = fo
+                    else:
+                        fo = open(path_stdout, 'a')
+                        fe = open(path_stderr, 'a')
+                    retcode = runner.write_stream_until_stop(
+                        fo, fe, flush_streams_each_time=True)
+                    fo.close()
+                    if path_stdout != path_stderr:
+                        fe.close()
                     return retcode
             else:
                 def _run_cmd(runner: ProcessRunner):

executor/engine/middle/capture.py

@@ -3,6 +3,10 @@
 from functools import update_wrapper
 from pathlib import Path
 import traceback
+import loguru
+
+
+LOGURU_HANDLERS = {}
 
 
 class Tee(object):
@@ -23,10 +27,18 @@ def flush(self):
 
     def __enter__(self):
         setattr(sys, self.stream_type, self)
+        fname = repr(self.file)
+        if fname not in LOGURU_HANDLERS:
+            loguru_handler = loguru.logger.add(self.file)
+            LOGURU_HANDLERS[fname] = loguru_handler
         return self
 
     def __exit__(self, _type, _value, _traceback):
         setattr(sys, self.stream_type, self.stream)
+        fname = repr(self.file)
+        if fname in LOGURU_HANDLERS:
+            loguru.logger.remove(LOGURU_HANDLERS[fname])
+            del LOGURU_HANDLERS[fname]
 
 
 class CaptureOut(object):
@@ -41,13 +53,26 @@ def __init__(
         self.capture_traceback = capture_traceback
 
     def __call__(self, *args, **kwargs) -> T.Any:
-        with open(self.stdout_file, 'w') as fo, \
-             open(self.stderr_file, 'w') as fe:
-            with Tee(fo, 'stdout'), Tee(fe, 'stderr'):
-                try:
-                    res = self.func(*args, **kwargs)
-                except Exception as e:
-                    if self.capture_traceback:
-                        traceback.print_exc(file=fe)
-                    raise e
+        if not self.stdout_file.parent.exists():
+            self.stdout_file.parent.mkdir(parents=True, exist_ok=True)  # pragma: no cover
+        if not self.stderr_file.parent.exists():
+            self.stderr_file.parent.mkdir(parents=True, exist_ok=True)  # pragma: no cover
+
+        if self.stdout_file == self.stderr_file:
+            outf = open(self.stdout_file, 'a')
+            errf = outf
+        else:
+            outf = open(self.stdout_file, 'a')
+            errf = open(self.stderr_file, 'a')
+        with Tee(outf, 'stdout'), Tee(errf, 'stderr'):
+            try:
+                res = self.func(*args, **kwargs)
+            except Exception as e:
+                if self.capture_traceback:
+                    traceback.print_exc()
+                raise e
+            finally:
+                outf.close()
+                if self.stdout_file != self.stderr_file:
+                    errf.close()
         return res

tests/test_job.py

@@ -1,4 +1,6 @@
 import asyncio
+import os
+import shutil
 import time
 
 import pytest
@@ -372,3 +374,27 @@ async def async_func(x):
             await engine.submit_async(job)
             await job.wait_until_status("done")
             assert job.result() == 2
+
+
+@pytest.mark.asyncio
+async def test_redirect_out_err():
+    with Engine() as engine:
+        def func():
+            import loguru
+            logger = loguru.logger
+            logger.info("test")
+            logger.error("test2")
+            import sys
+            print("test3", file=sys.stderr)
+            print("test4", file=sys.stdout)
+        unexist_dir = "tmp/1"
+        job = ProcessJob(func, redirect_out_err=f"{unexist_dir}/test.log")
+        await engine.submit_async(job)
+        await job.wait_until_status("done")
+        with open(f"{unexist_dir}/test.log", "r") as f:
+            content = f.read()
+            assert "test" in content
+            assert "test2" in content
+            assert "test3" in content
+            assert "test4" in content
+        shutil.rmtree(unexist_dir)

tests/test_subp_job.py

@@ -1,4 +1,6 @@
 import asyncio
+import pytest
+import shutil
 
 from executor.engine.core import Engine, EngineSetting
 from executor.engine.job.extend.subprocess import SubprocessJob
@@ -156,6 +158,20 @@ async def submit_job():
     asyncio.run(submit_job())
 
 
+@pytest.mark.asyncio
+async def test_redirect_out_err_str():
+    engine = Engine()
+    unexist_dir = "tmp/1"
+    job = SubprocessJob("python -c 'print(1 + 1); raise ValueError(\"error\")'", redirect_out_err=f"{unexist_dir}/test.log")
+    await engine.submit_async(job)
+    await job.join()
+    with open(f"{unexist_dir}/test.log") as f:
+        content = f.read()
+        assert "2" in content
+        assert "error" in content
+    #shutil.rmtree(unexist_dir)
+
+
 def test_cancel():
     engine = Engine()
 

tests/test_webapp_job.py

@@ -22,18 +22,18 @@ def test_run_webapp():
 
     async def submit_job():
         with pytest.raises(NotImplementedError):
-            job = WebappJob(run_simple_httpd, ip="2.2.2.2", port=8001, check_delta=0.5)
+            job = WebappJob(run_simple_httpd, ip="2.2.2.2", port=8088, check_delta=0.5)
 
         with pytest.raises(ValueError):
             job = WebappJob("python -m http.server -b ip port")
 
         with pytest.raises(TypeError):
             job = WebappJob(1)
 
-        job = WebappJob(run_simple_httpd, ip="127.0.0.1", port=8001, check_delta=0.5)
+        job = WebappJob(run_simple_httpd, ip="127.0.0.1", port=8088, check_delta=0.5)
         await engine.submit_async(job)
         await asyncio.sleep(5)
-        assert job.port == 8001
+        assert job.port == 8088
         assert job.status == "running"
         await job.cancel()
         assert job.status == "cancelled"