Skip to content

Commit 6e1d0fe

Browse files
feat: add host context support
Add the ability to pass "per invocation" host context accessible via CurrentPlugin. Additionally, allow referring to CurrentPlugin in type-inferred host functions.
1 parent 4e5051d commit 6e1d0fe

File tree

2 files changed

+54
-3
lines changed

2 files changed

+54
-3
lines changed

Diff for: extism/extism.py

+19-3
Original file line numberDiff line numberDiff line change
@@ -398,11 +398,18 @@ def __init__(self, namespace, name, func, user_data):
398398
arg_names = [arg for arg in hints.keys() if arg != "return"]
399399
returns = hints.pop("return", None)
400400

401+
uses_current_plugin = False
402+
if len(arg_names) > 0 and hints.get(arg_names[0], None) == CurrentPlugin:
403+
uses_current_plugin = True
404+
arg_names = arg_names[1:]
405+
401406
args = [_map_arg(arg, hints[arg]) for arg in arg_names]
407+
402408
returns = [] if returns is None else _map_ret(returns)
403409

404410
def inner_func(plugin, inputs, outputs, *user_data):
405-
inner_args = [
411+
first_arg = [plugin] if uses_current_plugin else []
412+
inner_args = first_arg + [
406413
extract(plugin, slot) for ((_, extract), slot) in zip(args, inputs)
407414
]
408415

@@ -523,6 +530,7 @@ def call(
523530
function_name: str,
524531
data: Union[str, bytes],
525532
parse: Callable[[Any], Any] = lambda xs: bytes(xs),
533+
host_context: Any = None,
526534
) -> Any:
527535
"""
528536
Call a function by name with the provided input data
@@ -533,11 +541,13 @@ def call(
533541
:raises: An :py:class:`extism.Error <.extism.Error>` if the guest function call was unsuccessful.
534542
:returns: The returned bytes from the guest function as interpreted by the ``parse`` parameter.
535543
"""
544+
545+
host_context = _ffi.new_handle(host_context)
536546
if isinstance(data, str):
537547
data = data.encode()
538548
self._check_error(
539-
_lib.extism_plugin_call(
540-
self.plugin, function_name.encode(), data, len(data)
549+
_lib.extism_plugin_call_with_host_context(
550+
self.plugin, function_name.encode(), data, len(data), host_context
541551
)
542552
)
543553
out_len = _lib.extism_plugin_output_length(self.plugin)
@@ -608,6 +618,12 @@ def memory(self, mem: Memory) -> _ffi.buffer:
608618
return None
609619
return _ffi.buffer(p + mem.offset, mem.length)
610620

621+
def host_context(self) -> Any:
622+
result = _lib.extism_current_plugin_host_context(self.pointer)
623+
if result == 0:
624+
return None
625+
return _ffi.from_handle(result)
626+
611627
def alloc(self, size: int) -> Memory:
612628
"""
613629
Allocate a new block of memory.

Diff for: tests/test_extism.py

+35
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
from collections import namedtuple
12
import unittest
23
import extism
34
import hashlib
@@ -148,6 +149,40 @@ def hello_world(
148149
self.assertIsInstance(result, Gribble)
149150
self.assertEqual(result.frobbitz(), "gromble robble")
150151

152+
def test_host_context(self):
153+
if not hasattr(typing, "Annotated"):
154+
return
155+
156+
# Testing two things here: one, if we see CurrentPlugin as the first arg, we pass it through.
157+
# Two, it's possible to refer to fetch the host context from the current plugin.
158+
@extism.host_fn(user_data=b"test")
159+
def hello_world(
160+
current_plugin: extism.CurrentPlugin,
161+
inp: typing.Annotated[dict, extism.Json],
162+
*user_data,
163+
) -> typing.Annotated[Gribble, extism.Pickle]:
164+
ctx = current_plugin.host_context()
165+
ctx.x = 1000
166+
return Gribble("robble")
167+
168+
plugin = extism.Plugin(
169+
self._manifest(functions=True), functions=[hello_world], wasi=True
170+
)
171+
172+
class Foo:
173+
x = 100
174+
y = 200
175+
176+
foo = Foo()
177+
178+
res = plugin.call("count_vowels", "aaa", host_context=foo)
179+
180+
self.assertEqual(foo.x, 1000)
181+
self.assertEqual(foo.y, 200)
182+
result = pickle.loads(res)
183+
self.assertIsInstance(result, Gribble)
184+
self.assertEqual(result.frobbitz(), "gromble robble")
185+
151186
def test_extism_plugin_cancel(self):
152187
plugin = extism.Plugin(self._loop_manifest())
153188
cancel_handle = plugin.cancel_handle()

0 commit comments

Comments
 (0)