Implement capturing stdout/stderr in Python (#311)

Bind a recent addition to the C API

Closes #34

Author: alexcrichton

tests/test_wasi.py

@@ -52,3 +52,65 @@ def preopen_nonexistent(self):
         config = WasiConfig()
         with self.assertRaises(WasmtimeError):
             config.preopen_dir('/path/to/nowhere', '/', DirPerms.READ_ONLY, FilePerms.READ_ONLY)
+
+    def test_custom_print(self):
+        linker = Linker(Engine())
+        linker.define_wasi()
+
+        stderr = ''
+        stdout = ''
+
+        def on_stdout(data: bytes) -> None:
+            nonlocal stdout
+            stdout += data.decode('utf8')
+
+        def on_stderr(data: bytes) -> None:
+            nonlocal stderr
+            stderr += data.decode('utf8')
+
+        module = Module(linker.engine, """
+            (module
+                (import "wasi_snapshot_preview1" "fd_write"
+                    (func $write (param i32 i32 i32 i32) (result i32)))
+
+                (memory (export "memory") 1)
+
+                (func $print
+                    (i32.store (i32.const 300) (i32.const 100)) ;; iov base
+                    (i32.store (i32.const 304) (i32.const 14))  ;; iov len
+
+                    (call $write
+                        (i32.const 1)    ;; fd 1 is stdout
+                        (i32.const 300)  ;; iovecs ptr
+                        (i32.const 1)    ;; iovecs len
+                        (i32.const 400)) ;; nwritten ptr
+                    if unreachable end   ;; verify no error
+
+                    (i32.store (i32.const 300) (i32.const 200)) ;; iov base
+                    (i32.store (i32.const 304) (i32.const 14))  ;; iov len
+
+                    (call $write
+                        (i32.const 2)    ;; fd 2 is stderr
+                        (i32.const 300)  ;; iovecs ptr
+                        (i32.const 1)    ;; iovecs len
+                        (i32.const 400)) ;; nwritten ptr
+                    if unreachable end   ;; verify no error
+                )
+
+                (start $print)
+
+                (data (i32.const 100) "Hello, stdout!")
+                (data (i32.const 200) "Hello, stderr!")
+            )
+        """)
+
+        wasi = WasiConfig()
+        wasi.stdout_custom = on_stdout
+        wasi.stderr_custom = on_stderr
+
+        store = Store(linker.engine)
+        store.set_wasi(wasi)
+        linker.instantiate(store, module)
+
+        self.assertEqual('Hello, stdout!', stdout)
+        self.assertEqual('Hello, stderr!', stderr)

wasmtime/_wasi.py

@@ -1,13 +1,15 @@
 import ctypes
-from ctypes import POINTER, c_char, c_char_p, cast
+import errno
+from ctypes import POINTER, c_char, c_char_p, cast, CFUNCTYPE, c_void_p
 from enum import Enum
 from os import PathLike
-from typing import Iterable, List, Union
+from typing import Iterable, List, Union, Callable
 
 from wasmtime import Managed, WasmtimeError
 
 from . import _ffi as ffi
 from ._config import setter_property
+from ._func import Slab
 
 
 def _encode_path(path: Union[str, bytes, PathLike]) -> bytes:
@@ -29,6 +31,11 @@ class FilePerms(Enum):
     WRITE_ONLY = ffi.wasi_file_perms_flags.WASMTIME_WASI_FILE_PERMS_WRITE.value
     READ_WRITE = ffi.wasi_file_perms_flags.WASMTIME_WASI_FILE_PERMS_READ.value | ffi.wasi_file_perms_flags.WASMTIME_WASI_FILE_PERMS_WRITE.value
 
+
+CustomOutput = Callable[[bytes], Union[int, None]]
+CUSTOM_OUTPUTS: Slab[CustomOutput] = Slab()
+
+
 class WasiConfig(Managed["ctypes._Pointer[ffi.wasi_config_t]"]):
 
     def __init__(self) -> None:
@@ -119,6 +126,15 @@ def stdout_file(self, path: str) -> None:
         if not res:
             raise WasmtimeError("failed to set stdout file")
 
+    @setter_property
+    def stdout_custom(self, callback: CustomOutput) -> None:
+        """
+        Sets a custom `callback` that is invoked whenever stdout is written to.
+        """
+        ffi.wasi_config_set_stdout_custom(
+            self.ptr(), custom_call,
+            CUSTOM_OUTPUTS.allocate(callback), custom_finalize)
+
     def inherit_stdout(self) -> None:
         """
         Configures this own process's stdout to be used as the WASI program's
@@ -145,6 +161,15 @@ def stderr_file(self, path: str) -> None:
         if not res:
             raise WasmtimeError("failed to set stderr file")
 
+    @setter_property
+    def stderr_custom(self, callback: CustomOutput) -> None:
+        """
+        Sets a custom `callback` that is invoked whenever stderr is written to.
+        """
+        ffi.wasi_config_set_stderr_custom(
+            self.ptr(), custom_call,
+            CUSTOM_OUTPUTS.allocate(callback), custom_finalize)
+
     def inherit_stderr(self) -> None:
         """
         Configures this own process's stderr to be used as the WASI program's
@@ -177,3 +202,24 @@ def to_char_array(strings: List[str]) -> "ctypes._Pointer[ctypes._Pointer[c_char
     for i, s in enumerate(strings):
         ptrs[i] = c_char_p(s.encode('utf-8'))
     return cast(ptrs, POINTER(POINTER(c_char)))
+
+
+@CFUNCTYPE(ctypes.c_ssize_t, c_void_p, POINTER(ctypes.c_ubyte), ctypes.c_size_t)
+def custom_call(idx, ptr, size): # type: ignore
+    try:
+        ty = ctypes.c_uint8 * size
+        arg = bytes(ty.from_address(ctypes.addressof(ptr.contents)))
+        ret = CUSTOM_OUTPUTS.get(idx or 0)(arg)
+        if ret is None:
+            return size
+        return ret
+    except Exception as e:
+        print('failed custom output, required to catch exception:', e)
+        return -errno.EIO
+
+
+@CFUNCTYPE(None, c_void_p)
+def custom_finalize(idx): # type: ignore
+    if CUSTOM_OUTPUTS:
+        CUSTOM_OUTPUTS.deallocate(idx or 0)
+    return None