feat: add support for env

Author: Niraj-Kamdar

packages/polywrap-client/polywrap_client/client.py

@@ -123,6 +123,9 @@ async def invoke(self, options: InvokerOptions) -> InvokeResult:
                 )
             ).unwrap()
 
+            env = self.get_env_by_uri(options.uri)
+            options.env = options.env or (env.env if env else None)
+
             result = await wrapper.invoke(options, invoker=self)
             if options.encode_result and not result.encoded:
                 encoded = msgpack_encode(result.result)

packages/polywrap-client/tests/cases/simple-env/wrap.info

@@ -1,7 +1,7 @@
 from pathlib import Path
 
 from polywrap_client import PolywrapClient
-from polywrap_core import Uri, InvokerOptions, InterfaceImplementations
+from polywrap_core import Uri, InvokerOptions, InterfaceImplementations, Env
 from polywrap_uri_resolvers import BaseUriResolver, SimpleFileReader
 
 from polywrap_client.client import PolywrapClientConfig
@@ -67,7 +67,32 @@ async def test_interface_implementation():
         f'fs/{Path(__file__).parent.joinpath("cases", "simple-interface", "wrapper").absolute()}'
     )
     args = {"arg": {"str": "hello", "uint8": 2}}
-    options = InvokerOptions(uri=uri, method="moduleMethod", args=args, encode_result=False)
+    options = InvokerOptions(
+        uri=uri, method="moduleMethod", args=args, encode_result=False
+    )
     result = await client.invoke(options)
 
     assert result.result == {"str": "hello", "uint8": 2}
+
+
+async def test_env():
+    uri_resolver = BaseUriResolver(
+        file_reader=SimpleFileReader(),
+        redirects={},
+    )
+
+    uri = Uri(f'fs/{Path(__file__).parent.joinpath("cases", "simple-env").absolute()}')
+    env = {"externalArray": [1, 2, 3], "externalString": "hello"}
+
+    client = PolywrapClient(
+        config=PolywrapClientConfig(
+            envs=[Env(uri=uri, env=env)],
+            resolver=uri_resolver,
+        )
+    )
+    options = InvokerOptions(
+        uri=uri, method="externalEnvMethod", args={}, encode_result=False
+    )
+    result = await client.invoke(options)
+
+    assert result.result == env

packages/polywrap-client/tests/cases/simple-env/wrap.wasm

@@ -79,6 +79,13 @@ def wrap_abort(
             f"__wrap_abort: {msg}\nFile: {file}\nLocation: [{line},{column}]"
         )
 
+    wrap_load_env_type = FuncType([ValType.i32()], [])
+
+    def wrap_load_env(ptr: int) -> None:
+        if not state.env:
+            raise WasmAbortError("env: is not set")
+        write_bytes(mem.data_ptr(store), mem.data_len(store), state.env, ptr)
+
     wrap_invoke_args_type = FuncType([ValType.i32(), ValType.i32()], [])
 
     def wrap_invoke_args(method_ptr: int, args_ptr: int) -> None:
@@ -309,6 +316,7 @@ def wrap_get_implementations_result(ptr: int) -> None:
     # TODO: use generics or any on wasmtime codebase to fix typings
     linker.define_func("wrap", "__wrap_debug_log", wrap_debug_log_type, wrap_debug_log)  # type: ignore partially unknown
     linker.define_func("wrap", "__wrap_abort", wrap_abort_type, wrap_abort)  # type: ignore partially unknown
+    linker.define_func("wrap", "__wrap_load_env", wrap_load_env_type, wrap_load_env)  # type: ignore partially unknown
 
     # invoke
     linker.define_func("wrap", "__wrap_invoke_args", wrap_invoke_args_type, wrap_invoke_args)  # type: ignore partially unknown