Merge pull request #96 from Peefy/feat-python-plugin

feat: python plugin in the kcl-lib

python/Makefile

@@ -7,7 +7,7 @@ fmt:
 	cargo fmt
 
 test:
-	python3 -m pip install pytest && python3 -m pytest
+	export PYTHONPATH=${PWD} && python3 -m pip install pytest && python3 -m pytest
 
 deps:
 	python3 -m venv venv

python/kcl_lib/api/service.py

@@ -97,9 +97,11 @@ def call(self, name: str, args):
         args_serialized = args.SerializeToString()
 
         # Call the service function and get the result
-        result = bytes(kcl_lib.call_with_plugin_agent(
-            name.encode("utf-8"), args_serialized, self.plugin_agent
-        ))
+        result = bytes(
+            kcl_lib.call_with_plugin_agent(
+                name.encode("utf-8"), args_serialized, self.plugin_agent
+            )
+        )
         if result.startswith(b"ERROR"):
             raise Exception(str(result))
         msg = self.create_method_resp_message(name)

python/kcl_lib/plugin/__init__.py

@@ -0,0 +1,3 @@
+from .plugin import Plugin, register_plugin, plugin_agent_addr
+
+__all__ = ["Plugin", "register_plugin", "plugin_agent_addr"]

python/kcl_lib/plugin/plugin.py

@@ -0,0 +1,82 @@
+import json
+import typing
+from ctypes import *
+from dataclasses import dataclass
+
+
+@dataclass
+class Plugin:
+    name: str
+    method_map: typing.Dict[str, typing.Callable]
+
+
+class PluginContext:
+    def __init__(self):
+        self._plugin_dict: typing.Dict[str, Plugin] = {}
+
+    def call_method(self, name: str, args_json: str, kwargs_json: str) -> str:
+        return self._call_py_method(name, args_json, kwargs_json)
+
+    def _call_py_method(self, name: str, args_json: str, kwargs_json: str) -> str:
+        try:
+            return self._call_py_method_unsafe(name, args_json, kwargs_json)
+        except Exception as e:
+            return json.dumps({"__kcl_PanicInfo__": f"{e}"})
+
+    def _call_py_method_unsafe(
+        self, name: str, args_json: str, kwargs_json: str
+    ) -> str:
+        dotIdx = name.rfind(".")
+        if dotIdx < 0:
+            return ""
+
+        # kcl_plugin.<module_path>.<method_name>
+        module_path = name[:dotIdx]
+        method_name = name[dotIdx + 1 :]
+        plugin_name = module_path[module_path.rfind(".") + 1 :]
+
+        method_func = self._plugin_dict.get(
+            plugin_name, Plugin(name="", method_map={})
+        ).method_map.get(method_name, None)
+
+        args = []
+        kwargs = {}
+
+        if args_json:
+            args = json.loads(args_json)
+            if not isinstance(args, list):
+                return ""
+
+        if kwargs_json:
+            kwargs = json.loads(kwargs_json)
+            if not isinstance(kwargs, dict):
+                return ""
+
+        result = method_func(*args, **kwargs) if method_func else None
+        return json.dumps(result)
+
+
+__plugin_context__ = PluginContext()
+__plugin_method_agent_buffer__ = create_string_buffer(1024)
+
+
+def register_plugin(
+    name: str, method_map: typing.Dict[str, typing.Callable]
+) -> c_char_p:
+    __plugin_context__._plugin_dict[name] = Plugin(name, method_map)
+
+
+@CFUNCTYPE(c_char_p, c_char_p, c_char_p, c_char_p)
+def plugin_method_agent(method: str, args_json: str, kwargs_json: str) -> c_char_p:
+    method = str(method, encoding="utf-8")
+    args_json = str(args_json, encoding="utf-8")
+    kwargs_json = str(kwargs_json, encoding="utf-8")
+
+    json_result = __plugin_context__.call_method(method, args_json, kwargs_json)
+
+    global __plugin_method_agent_buffer__
+    __plugin_method_agent_buffer__ = create_string_buffer(json_result.encode("utf-8"))
+    return addressof(__plugin_method_agent_buffer__)
+
+
+plugin_agent_addr = cast(plugin_method_agent, c_void_p).value

python/tests/plugin_test.py

@@ -0,0 +1,12 @@
+import kcl_lib.plugin as plugin
+import kcl_lib.api as api
+
+TEST_FILE = "./tests/test_data/plugin.k"
+
+
+def test_exec_api():
+    plugin.register_plugin("my_plugin", {"add": lambda x, y: x + y})
+    result = api.API(plugin.plugin_agent_addr).exec_program(
+        api.ExecProgram_Args(k_filename_list=[TEST_FILE])
+    )
+    assert result.yaml_result == "result: 2"

python/tests/test_data/plugin.k

@@ -0,0 +1,3 @@
+import kcl_plugin.my_plugin
+
+result = my_plugin.add(1, 1)