Skip to content

Commit

Permalink
feat: add host context support
Browse files Browse the repository at this point in the history
Add the ability to pass "per invocation" host context accessible via
CurrentPlugin. Additionally, allow referring to CurrentPlugin in type-inferred
host functions.
  • Loading branch information
chrisdickinson committed Nov 22, 2024
1 parent 4e5051d commit 6e1d0fe
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 3 deletions.
22 changes: 19 additions & 3 deletions extism/extism.py
Original file line number Diff line number Diff line change
Expand Up @@ -398,11 +398,18 @@ def __init__(self, namespace, name, func, user_data):
arg_names = [arg for arg in hints.keys() if arg != "return"]
returns = hints.pop("return", None)

uses_current_plugin = False
if len(arg_names) > 0 and hints.get(arg_names[0], None) == CurrentPlugin:
uses_current_plugin = True
arg_names = arg_names[1:]

args = [_map_arg(arg, hints[arg]) for arg in arg_names]

returns = [] if returns is None else _map_ret(returns)

def inner_func(plugin, inputs, outputs, *user_data):
inner_args = [
first_arg = [plugin] if uses_current_plugin else []
inner_args = first_arg + [
extract(plugin, slot) for ((_, extract), slot) in zip(args, inputs)
]

Expand Down Expand Up @@ -523,6 +530,7 @@ def call(
function_name: str,
data: Union[str, bytes],
parse: Callable[[Any], Any] = lambda xs: bytes(xs),
host_context: Any = None,
) -> Any:
"""
Call a function by name with the provided input data
Expand All @@ -533,11 +541,13 @@ def call(
:raises: An :py:class:`extism.Error <.extism.Error>` if the guest function call was unsuccessful.
:returns: The returned bytes from the guest function as interpreted by the ``parse`` parameter.
"""

host_context = _ffi.new_handle(host_context)
if isinstance(data, str):
data = data.encode()
self._check_error(
_lib.extism_plugin_call(
self.plugin, function_name.encode(), data, len(data)
_lib.extism_plugin_call_with_host_context(
self.plugin, function_name.encode(), data, len(data), host_context
)
)
out_len = _lib.extism_plugin_output_length(self.plugin)
Expand Down Expand Up @@ -608,6 +618,12 @@ def memory(self, mem: Memory) -> _ffi.buffer:
return None
return _ffi.buffer(p + mem.offset, mem.length)

def host_context(self) -> Any:
result = _lib.extism_current_plugin_host_context(self.pointer)
if result == 0:
return None
return _ffi.from_handle(result)

def alloc(self, size: int) -> Memory:
"""
Allocate a new block of memory.
Expand Down
35 changes: 35 additions & 0 deletions tests/test_extism.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from collections import namedtuple
import unittest
import extism
import hashlib
Expand Down Expand Up @@ -148,6 +149,40 @@ def hello_world(
self.assertIsInstance(result, Gribble)
self.assertEqual(result.frobbitz(), "gromble robble")

def test_host_context(self):
if not hasattr(typing, "Annotated"):
return

# Testing two things here: one, if we see CurrentPlugin as the first arg, we pass it through.
# Two, it's possible to refer to fetch the host context from the current plugin.
@extism.host_fn(user_data=b"test")
def hello_world(
current_plugin: extism.CurrentPlugin,
inp: typing.Annotated[dict, extism.Json],
*user_data,
) -> typing.Annotated[Gribble, extism.Pickle]:
ctx = current_plugin.host_context()
ctx.x = 1000
return Gribble("robble")

plugin = extism.Plugin(
self._manifest(functions=True), functions=[hello_world], wasi=True
)

class Foo:
x = 100
y = 200

foo = Foo()

res = plugin.call("count_vowels", "aaa", host_context=foo)

self.assertEqual(foo.x, 1000)
self.assertEqual(foo.y, 200)
result = pickle.loads(res)
self.assertIsInstance(result, Gribble)
self.assertEqual(result.frobbitz(), "gromble robble")

def test_extism_plugin_cancel(self):
plugin = extism.Plugin(self._loop_manifest())
cancel_handle = plugin.cancel_handle()
Expand Down

0 comments on commit 6e1d0fe

Please sign in to comment.