feat: add custom logging

example.py

@@ -5,11 +5,9 @@
 import hashlib
 import pathlib
 
-from extism import Function, host_fn, ValType, Plugin, set_log_file, Json
+from extism import Function, host_fn, ValType, Plugin, set_log_custom, Json
 from typing import Annotated
 
-set_log_file("stderr", "trace")
-
 
 @host_fn(user_data=b"Hello again!")
 def hello_world(inp: Annotated[dict, Json], *a_string) -> Annotated[dict, Json]:
@@ -25,7 +23,8 @@ def count_vowels(data):
 
 
 def main(args):
-    set_log_file("stderr", "trace")
+    logs = []
+    logBuffer = set_log_custom(lambda s: logs.append(s.strip()), "trace")
     if len(args) > 1:
         data = args[1].encode()
     else:
@@ -48,6 +47,12 @@ def main(args):
     assert j["count"] == count_vowels(data)
     assert j["roundtrip"] == 1
 
+    # Drain logs and print
+    logBuffer.drain()
+    print("Dumping logs", len(logs))
+    for line in logs:
+        print(line)
+
 
 if __name__ == "__main__":
     main(sys.argv)

extism/__init__.py

@@ -7,6 +7,7 @@
     Error,
     Plugin,
     set_log_file,
+    set_log_custom,
     extism_version,
     host_fn,
     Function,
@@ -24,6 +25,7 @@
     "Error",
     "CurrentPlugin",
     "set_log_file",
+    "set_log_custom",
     "extism_version",
     "Memory",
     "host_fn",

extism/extism.py

@@ -163,6 +163,45 @@ def set_log_file(
     _lib.extism_log_file(file.encode(), c_level)
 
 
+class CustomLogger:
+    def __init__(self, f):
+        self.callback = None
+        self.set_callback(f)
+
+    def set_callback(self, f):
+        @_ffi.callback("void(char*, unsigned long)")
+        def callback(ptr, len):
+            f(_ffi.string(ptr, len).decode())
+
+        self.callback = callback
+
+    def drain(self):
+        if self.callback is not None:
+            _lib.extism_log_drain(self.callback)
+
+    def __del__(self):
+        self.drain()
+
+
+def set_log_custom(
+    f, level: Optional[Literal["debug", "error", "trace", "warn"]] = None
+):
+    """
+    Enables buffered logging, this is a global configuration
+
+    :param f: The callback function, takes a string argument and no return value.
+    :param level: The debug level.
+
+    :returns: a CustomLogger with a `drain` method that can be used to handle the buffered logs.
+    """
+    c_level = level or _ffi.NULL
+    if isinstance(level, str):
+        c_level = level.encode()
+
+    _lib.extism_log_custom(c_level)
+    return CustomLogger(f)
+
+
 def extism_version() -> str:
     """
     Gets the Extism version string