Merge pull request #223 from xen0n/cmd-plugin

Support plugin-provided Ruyi subcommands

ruyi/cli/builtin_commands.py

@@ -1,5 +1,6 @@
 from ..device import provision_cli as provision_cli
 from ..mux.venv import venv_cli as venv_cli
+from ..pluginhost import plugin_cli as plugin_cli
 from ..ruyipkg import admin_cli as admin_cli
 from ..ruyipkg import news_cli as news_cli
 from ..ruyipkg import pkg_cli as pkg_cli

ruyi/cli/cmd.py

@@ -163,3 +163,15 @@ def configure_args(cls, p: argparse.ArgumentParser) -> None:
             action="store_true",
             help="Give the output in a machine-friendly format if applicable",
         )
+
+
+# Repo admin commands
+class AdminCommand(
+    RootCommand,
+    cmd="admin",
+    has_subcommands=True,
+    # https://github.com/python/cpython/issues/67037
+    # help=argparse.SUPPRESS,
+    help="(NOT FOR REGULAR USERS) Subcommands for managing Ruyi repos",
+):
+    pass

ruyi/pluginhost/__init__.py

@@ -57,33 +57,42 @@ def __init__(
     @abc.abstractmethod
     def make_loader(
         self,
-        plugin_root: pathlib.Path,
         originating_file: pathlib.Path,
         module_cache: MutableMapping[str, ModuleTy],
+        is_cmd: bool,
     ) -> "BasePluginLoader[ModuleTy]":
         raise NotImplementedError
 
     @abc.abstractmethod
     def make_evaluator(self) -> EvalTy:
         raise NotImplementedError
 
-    def load_plugin(self, plugin_id: str) -> None:
+    @property
+    def plugin_root(self) -> pathlib.Path:
+        return self._plugin_root
+
+    def load_plugin(self, plugin_id: str, is_cmd: bool) -> None:
         plugin_dir = paths.get_plugin_dir(plugin_id, self._plugin_root)
 
         loader = self.make_loader(
-            self._plugin_root,
             plugin_dir / paths.PLUGIN_ENTRYPOINT_FILENAME,
             self._module_cache,
+            is_cmd,
         )
         loaded_plugin = loader.load_this_plugin()
         self._loaded_plugins[plugin_id] = loaded_plugin
 
     def is_plugin_loaded(self, plugin_id: str) -> bool:
         return plugin_id in self._loaded_plugins
 
-    def get_from_plugin(self, plugin_id: str, key: str) -> object | None:
+    def get_from_plugin(
+        self,
+        plugin_id: str,
+        key: str,
+        is_cmd_plugin: bool = False,
+    ) -> object | None:
         if not self.is_plugin_loaded(plugin_id):
-            self.load_plugin(plugin_id)
+            self.load_plugin(plugin_id, is_cmd_plugin)
 
         if plugin_id not in self._value_cache:
             self._value_cache[plugin_id] = {}
@@ -111,19 +120,26 @@ class BasePluginLoader(Generic[ModuleTy], metaclass=abc.ABCMeta):
 
     def __init__(
         self,
-        root: pathlib.Path,
+        phctx: PluginHostContext[ModuleTy, SupportsEvalFunction],
         originating_file: pathlib.Path,
         module_cache: MutableMapping[str, ModuleTy],
+        is_cmd: bool,
     ) -> None:
-        self.root = root
+        self._phctx = phctx
         self.originating_file = originating_file
         self.module_cache = module_cache
+        self.is_cmd = is_cmd
+
+    @property
+    def root(self) -> pathlib.Path:
+        return self._phctx.plugin_root
 
     def make_sub_loader(self, originating_file: pathlib.Path) -> Self:
         return self.__class__(
-            self.root,
+            self._phctx,
             originating_file,
             self.module_cache,
+            self.is_cmd,
         )
 
     def load_this_plugin(self) -> ModuleTy:
@@ -142,6 +158,7 @@ def _load(self, path: str, is_root: bool) -> ModuleTy:
                 self.root,
                 False,
                 self.originating_file,
+                self.is_cmd,
             )
         resolved_path_str = str(resolved_path)
         if resolved_path_str in self.module_cache:
@@ -151,9 +168,10 @@ def _load(self, path: str, is_root: bool) -> ModuleTy:
         plugin_dir = self.root / plugin_id
 
         host_bridge = api.make_ruyi_plugin_api_for_module(
-            self.root,
+            self._phctx,
             resolved_path,
             plugin_dir,
+            self.is_cmd,
         )
 
         mod = self.do_load_module(

ruyi/pluginhost/api.py

@@ -1,25 +1,35 @@
+from contextlib import AbstractContextManager
 import pathlib
 import subprocess
 import time
 import tomllib
-from typing import Any, Callable
+from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast
 
-from ruyi import log
-from ruyi.cli import user_input
-from ruyi.version import RUYI_SEMVER
+from .. import log
+from ..cli import user_input
+from ..version import RUYI_SEMVER
 from .paths import resolve_ruyi_load_path
 
+if TYPE_CHECKING:
+    from . import PluginHostContext, SupportsEvalFunction, SupportsGetOption
+
+T = TypeVar("T")
+U = TypeVar("U")
+
 
 class RuyiHostAPI:
     def __init__(
         self,
-        plugin_root: pathlib.Path,
+        phctx: "PluginHostContext[SupportsGetOption, SupportsEvalFunction]",
         this_file: pathlib.Path,
         this_plugin_dir: pathlib.Path,
+        allow_host_fs_access: bool,
     ) -> None:
-        self._plugin_root = plugin_root
+        self._phctx = phctx
         self._this_file = this_file
         self._this_plugin_dir = this_plugin_dir
+        self._ev = phctx.make_evaluator()
+        self._allow_host_fs_access = allow_host_fs_access
 
         self._logger = RuyiPluginLogger()
 
@@ -34,9 +44,10 @@ def ruyi_plugin_api_rev(self) -> int:
     def load_toml(self, path: str) -> object:
         resolved_path = resolve_ruyi_load_path(
             path,
-            self._plugin_root,
+            self._phctx.plugin_root,
             True,
             self._this_file,
+            self._allow_host_fs_access,
         )
         with open(resolved_path, "rb") as f:
             return tomllib.load(f)
@@ -71,6 +82,14 @@ def call_subprocess_argv(
     def sleep(self, seconds: float, /) -> None:
         return time.sleep(seconds)
 
+    def with_(
+        self,
+        cm: AbstractContextManager[T],
+        fn: object | Callable[[T], U],
+    ) -> U:
+        with cm as obj:
+            return cast(U, self._ev.eval_function(fn, obj))
+
 
 class RuyiPluginLogger:
     def __init__(self) -> None:
@@ -123,9 +142,10 @@ def F(
 
 
 def _ruyi_plugin_rev(
-    plugin_root: pathlib.Path,
+    phctx: "PluginHostContext[SupportsGetOption, SupportsEvalFunction]",
     this_file: pathlib.Path,
     this_plugin_dir: pathlib.Path,
+    allow_host_fs_access: bool,
     rev: object,
 ) -> RuyiHostAPI:
     if not isinstance(rev, int):
@@ -134,12 +154,27 @@ def _ruyi_plugin_rev(
         raise ValueError(
             f"Ruyi plugin API revision {rev} is not supported by this Ruyi"
         )
-    return RuyiHostAPI(plugin_root, this_file, this_plugin_dir)
+    return RuyiHostAPI(
+        phctx,
+        this_file,
+        this_plugin_dir,
+        allow_host_fs_access,
+    )
 
 
 def make_ruyi_plugin_api_for_module(
-    plugin_root: pathlib.Path,
+    phctx: "PluginHostContext[SupportsGetOption, SupportsEvalFunction]",
     this_file: pathlib.Path,
     this_plugin_dir: pathlib.Path,
+    is_cmd: bool,
 ) -> Callable[[object], RuyiHostAPI]:
-    return lambda rev: _ruyi_plugin_rev(plugin_root, this_file, this_plugin_dir, rev)
+    # Only allow access to host FS when we're being loaded as a command plugin
+    allow_host_fs_access = is_cmd
+
+    return lambda rev: _ruyi_plugin_rev(
+        phctx,
+        this_file,
+        this_plugin_dir,
+        allow_host_fs_access,
+        rev,
+    )

ruyi/pluginhost/paths.py

@@ -22,6 +22,7 @@ def resolve_ruyi_load_path(
     plugin_root: pathlib.Path,
     is_for_data: bool,
     originating_file: pathlib.Path,
+    allow_host_fs_access: bool,
 ) -> pathlib.Path:
     parsed = urlparse(path)
     if parsed.params or parsed.query or parsed.fragment:
@@ -80,6 +81,22 @@ def resolve_ruyi_load_path(
                 plugin_id=parsed.netloc,
             )
 
+        case "host":
+            if not allow_host_fs_access:
+                raise RuntimeError("the host protocol is not allowed in this context")
+
+            if not parsed.path:
+                raise RuntimeError(
+                    "empty path segment is not allowed for host:// load paths"
+                )
+
+            if parsed.netloc:
+                raise RuntimeError(
+                    "non-empty location is not allowed for host:// load paths"
+                )
+
+            return pathlib.Path(parsed.path)
+
         case _:
             raise RuntimeError(
                 f"unsupported Ruyi Starlark load path scheme {parsed.scheme}"

ruyi/pluginhost/plugin_cli.py

@@ -0,0 +1,35 @@
+import argparse
+
+from ..cli.cmd import AdminCommand
+from ..config import GlobalConfig
+from ..ruyipkg.repo import MetadataRepo
+
+
+class AdminRunPluginCommand(
+    AdminCommand,
+    cmd="run-plugin-cmd",
+    help="Run a plugin-defined command",
+):
+    @classmethod
+    def configure_args(cls, p: argparse.ArgumentParser) -> None:
+        p.add_argument(
+            "cmd_name",
+            type=str,
+            metavar="COMMAND-NAME",
+            help="Command name",
+        )
+        p.add_argument(
+            "cmd_args",
+            type=str,
+            nargs="*",
+            metavar="COMMAND-ARG",
+            help="Arguments to pass to the plugin command",
+        )
+
+    @classmethod
+    def main(cls, cfg: GlobalConfig, args: argparse.Namespace) -> int:
+        cmd_name = args.cmd_name
+        cmd_args = args.cmd_args
+
+        mr = MetadataRepo(cfg)
+        return mr.run_plugin_cmd(cmd_name, cmd_args)

ruyi/pluginhost/unsandboxed.py

@@ -72,11 +72,11 @@ class UnsandboxedPluginHostContext(
 ):
     def make_loader(
         self,
-        plugin_root: pathlib.Path,
         originating_file: pathlib.Path,
         module_cache: MutableMapping[str, UnsandboxedModuleDict],
+        is_cmd: bool,
     ) -> BasePluginLoader[UnsandboxedModuleDict]:
-        return UnsandboxedRuyiPluginLoader(plugin_root, originating_file, module_cache)
+        return UnsandboxedRuyiPluginLoader(self, originating_file, module_cache, is_cmd)
 
     def make_evaluator(self) -> UnsandboxedTrivialEvaluator:
         return UnsandboxedTrivialEvaluator()

ruyi/ruyipkg/admin_cli.py

@@ -10,25 +10,13 @@
 from tomlkit.items import AoT, Table
 
 from .. import log
-from ..cli.cmd import RootCommand
+from ..cli.cmd import AdminCommand
 from ..config import GlobalConfig
 from . import checksum
 from .canonical_dump import dump_canonical_package_manifest_toml
 from .pkg_manifest import DistfileDeclType, PackageManifest, RestrictKind
 
 
-# Repo admin commands
-class AdminCommand(
-    RootCommand,
-    cmd="admin",
-    has_subcommands=True,
-    # https://github.com/python/cpython/issues/67037
-    # help=argparse.SUPPRESS,
-    help="(NOT FOR REGULAR USERS) Subcommands for managing Ruyi repos",
-):
-    pass
-
-
 class AdminManifestCommand(
     AdminCommand,
     cmd="manifest",

ruyi/ruyipkg/repo.py

@@ -499,3 +499,23 @@ def get_provisioner_config(self) -> ProvisionerConfig | None:
             self.ensure_provisioner_config_cache()
         assert self._provisioner_config_cache is not None
         return self._provisioner_config_cache[0]
+
+    def run_plugin_cmd(self, cmd_name: str, args: list[str]) -> int:
+        plugin_id = f"ruyi-cmd-{cmd_name.lower()}"
+
+        plugin_entrypoint = self._plugin_host_ctx.get_from_plugin(
+            plugin_id,
+            "plugin_cmd_main_v1",
+            is_cmd_plugin=True,  # allow access to host FS for command plugins
+        )
+        if plugin_entrypoint is None:
+            raise RuntimeError(f"cmd entrypoint not found in plugin '{plugin_id}'")
+
+        ret = self.eval_plugin_fn(plugin_entrypoint, args)
+        if not isinstance(ret, int):
+            log.W(
+                f"unexpected return type of cmd plugin '{plugin_id}': {type(ret)} is not int."
+            )
+            log.I("forcing return code to 1; the plugin should be fixed")
+            ret = 1
+        return ret

tests/fixtures/__init__.py

@@ -12,6 +12,10 @@ def __init__(self, module: resources.Package | None = None) -> None:
     def path(self, *frags: str) -> AbstractContextManager[pathlib.Path]:
         return resources.as_file(resources.files(self.module).joinpath(*frags))
 
+    def plugin_suite(self, suite_name: str) -> AbstractContextManager[pathlib.Path]:
+        path = resources.files(self.module).joinpath("plugins_suites", suite_name)
+        return resources.as_file(path)
+
 
 @pytest.fixture
 def ruyi_file() -> RuyiFileFixtureFactory:

tests/fixtures/plugins_suites/with_/foo/mod.star

@@ -0,0 +1,17 @@
+RUYI = ruyi_plugin_rev(1)
+
+
+def fn1(mgr):
+    def _with_inner(obj):
+        return obj * 2
+    return RUYI.with_(mgr, _with_inner)
+
+
+def fn2(mgr):
+    def _with_inner(obj):
+        return mgr.NoNeXiStEnT
+    return RUYI.with_(mgr, _with_inner)
+
+
+def fn3(mgr, py_fn):
+    return RUYI.with_(mgr, py_fn)