Implemented way to fetch the raw bytes from a File in Python

edomora97 · edomora97 · commit 9e545e88b39e · 2018-12-30T16:03:36.000+01:00
diff --git a/cpp/frontend/python/frontend.cpp b/cpp/frontend/python/frontend.cpp
@@ -118,6 +118,20 @@ PYBIND11_MODULE(task_maker_frontend, m) {
                  });
            },
            "callback"_a)
+      .def("getContentsAsBytes",
+           [](frontend::File &f, std::function<void(pybind11::bytes)> cb) {
+             f.getContentsAsString(
+                 [cb = destroy_with_gil(cb)](std::string s) mutable {
+                   pybind11::gil_scoped_acquire acquire;
+                   try {
+                     (*cb)(s);
+                   } catch (pybind11::error_already_set &exc) {
+                     std::cerr << __FILE__ << ":" << __LINE__ << " "
+                               << exc.what() << std::endl;
+                     _Exit(1);
+                   }
+                 });
+           })
       .def("getContentsToFile", &frontend::File::getContentsToFile, "path"_a,
            "overwrite"_a = true, "exist_ok"_a = true);
 
diff --git a/python/remote/__init__.py b/python/remote/__init__.py
@@ -47,14 +47,20 @@ def start(self):
         """
 
         def stop_server(_1: int, _2: Any) -> None:
-            self.frontend.stopEvaluation()
-            self.stopped = True
+            self.stop()
 
         signal.signal(signal.SIGINT, stop_server)
         signal.signal(signal.SIGTERM, stop_server)
 
         self.frontend.evaluate()
 
+    def stop(self):
+        """
+        Stop the current evaluation
+        """
+        self.frontend.stopEvaluation()
+        self.stopped = True
+
 
 class Execution:
     """
@@ -82,6 +88,8 @@ def __init__(self,
                  stderr_fifo: Union[Fifo, None] = None,
                  store_stdout: bool = False,
                  store_stderr: bool = False,
+                 store_stdout_bytes: bool = False,
+                 store_stderr_bytes: bool = False,
                  inputs: Dict[str, File] = None,
                  outputs: Iterable[Union[str, Tuple[str, bool]]] = ()):
         """
@@ -128,11 +136,15 @@ def __init__(self,
         self.stderr_fifo = stderr_fifo
         self.store_stdout = store_stdout
         self.store_stderr = store_stderr
+        self.store_stdout_bytes = store_stdout_bytes
+        self.store_stderr_bytes = store_stderr_bytes
 
         self.stdout = None  # type: Optional[File]
         self.stderr = None  # type: Optional[File]
         self._stdout = None  # type: Optional[str]
         self._stderr = None  # type: Optional[str]
+        self._stdout_bytes = None  # type: Optional[bytes]
+        self._stderr_bytes = None  # type: Optional[bytes]
         self._outputs = dict()  # type: Dict[str, File]
         if not inputs:
             inputs = dict()
@@ -233,6 +245,10 @@ def __init__(self,
             self.stdout.getContentsAsString(self._get_stdout_internal)
         if self.store_stderr:
             self.stderr.getContentsAsString(self._get_stderr_internal)
+        if self.store_stdout_bytes:
+            self.stdout.getContentsAsBytes(self._get_stdout_bytes_internal)
+        if self.store_stderr_bytes:
+            self.stderr.getContentsAsBytes(self._get_stderr_bytes_internal)
 
         self._execution.notifyStart(self._notify_start_internal)
         # getResult should be the last thing done on _execution
@@ -275,13 +291,25 @@ def _get_stderr_internal(self, stderr: str):
         self._stderr = stderr
         self._on_done_internal()
 
+    def _get_stdout_bytes_internal(self, stdout: bytes):
+        self._stdout_bytes = stdout
+        self._on_done_internal()
+
+    def _get_stderr_bytes_internal(self, stderr: bytes):
+        self._stderr_bytes = stderr
+        self._on_done_internal()
+
     def _on_done_internal(self):
         if not self._on_done_cb:
             return
         if self.store_stdout and self._stdout is None:
             return
         if self.store_stderr and self._stderr is None:
             return
+        if self.store_stdout_bytes and self._stdout_bytes is None:
+            return
+        if self.store_stderr_bytes and self._stderr_bytes is None:
+            return
         if self.result is None:
             return
         self._on_done_cb(self.result)
@@ -312,6 +340,32 @@ def stderr_content(self) -> str:
                                "the execution has completed")
         return self._stderr
 
+    @property
+    def stdout_content_bytes(self) -> bytes:
+        """
+        The content of stdout as bytes, must be called after the execution has
+        completed and only if store_stdout_bytes has been set to True
+        """
+        if not self.store_stdout_bytes:
+            raise ValueError("Stdout was not captured as bytes")
+        if not self.result:
+            raise RuntimeError("stdout_content_bytes must be called after "
+                               "the execution has completed")
+        return self._stdout_bytes
+
+    @property
+    def stderr_content_bytes(self) -> bytes:
+        """
+        The content of stderr as bytes, must be called after the execution has
+        completed and only if store_stderr_bytes has been set to True
+        """
+        if not self.store_stderr_bytes:
+            raise ValueError("Stderr was not captured as bytes")
+        if not self.result:
+            raise RuntimeError("stderr_content_bytes must be called after "
+                               "the execution has completed")
+        return self._stderr_bytes
+
     def output(self, path: str) -> File:
         """
         The file generated by the execution on that path. Must be called after
diff --git a/python/source_file.py b/python/source_file.py
@@ -88,6 +88,17 @@ def __init__(self, path: str, unit_name: str, exe_name: str,
     def prepared(self) -> bool:
         return self.pool is not None
 
+    def unprepare(self):
+        """
+        Unprepare the source file. Useful for recompiling it in a different
+        execution pool.
+        """
+        self.pool = None
+        self.executable = None
+        self.compilation = None
+        self.compilation_stderr = None
+        self.compilation_stdout = None
+
     def prepare(self, pool: ExecutionPool):
         """
         Prepare the source file for execution, compile the source if needed.
