Privatize APIs by renaming or removing them from __all__ (#6021)

redbot/core/_rpc.py

@@ -0,0 +1,199 @@
+import asyncio
+import sys
+from typing import Optional
+
+from aiohttp import web
+from aiohttp_json_rpc import JsonRpc
+from aiohttp_json_rpc.rpc import JsonRpcMethod
+
+import logging
+
+from redbot.core._cli import ExitCodes
+
+log = logging.getLogger("red.rpc")
+
+__all__ = ("RPC", "RPCMixin", "get_name")
+
+
+def get_name(func, prefix=""):
+    class_name = prefix or func.__self__.__class__.__name__.lower()
+    func_name = func.__name__.strip("_")
+    if class_name == "redrpc":
+        return func_name.upper()
+    return f"{class_name}__{func_name}".upper()
+
+
+class RedRpc(JsonRpc):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.add_methods(("", self.get_method_info))
+
+    def _add_method(self, method, name="", prefix=""):
+        if not asyncio.iscoroutinefunction(method):
+            return
+
+        name = name or get_name(method, prefix)
+
+        self.methods[name] = JsonRpcMethod(method)
+
+    def remove_method(self, method):
+        meth_name = get_name(method)
+        new_methods = {}
+        for name, meth in self.methods.items():
+            if name != meth_name:
+                new_methods[name] = meth
+        self.methods = new_methods
+
+    def remove_methods(self, prefix: str):
+        new_methods = {}
+        for name, meth in self.methods.items():
+            splitted = name.split("__")
+            if len(splitted) < 2 or splitted[0] != prefix:
+                new_methods[name] = meth
+        self.methods = new_methods
+
+    async def get_method_info(self, request):
+        method_name = request.params[0]
+        if method_name in self.methods:
+            return self.methods[method_name].__doc__
+        return "No docstring available."
+
+
+class RPC:
+    """
+    RPC server manager.
+    """
+
+    app: web.Application
+    _rpc: RedRpc
+    _runner: web.AppRunner
+
+    def __init__(self):
+        self._site: Optional[web.TCPSite] = None
+        self._started = False
+
+    async def _pre_login(self) -> None:
+        self.app = web.Application()
+        self._rpc = RedRpc()
+        self.app.router.add_route("*", "/", self._rpc.handle_request)
+
+        self._runner = web.AppRunner(self.app)
+
+    async def initialize(self, port: int):
+        """
+        Finalizes the initialization of the RPC server and allows it to begin
+        accepting queries.
+        """
+        try:
+            # This ensures self._started can't be assigned
+            # except with both other functions
+            # and isn't subject to a really really stupid but complex
+            # issue on windows with catching specific
+            # exceptions related to shutdown conditions in asyncio applications.
+            self._started, _discard, self._site = (
+                True,
+                await self._runner.setup(),
+                web.TCPSite(self._runner, host="127.0.0.1", port=port, shutdown_timeout=0),
+            )
+        except Exception as exc:
+            log.exception("RPC setup failure", exc_info=exc)
+            sys.exit(ExitCodes.CRITICAL)
+        else:
+            await self._site.start()
+            log.debug("Created RPC server listener on port %s", port)
+
+    async def close(self):
+        """
+        Closes the RPC server.
+        """
+        if self._started:
+            await self.app.shutdown()
+            await self._runner.cleanup()
+
+    def add_method(self, method, prefix: str = None):
+        if prefix is None:
+            prefix = method.__self__.__class__.__name__.lower()
+
+        if not asyncio.iscoroutinefunction(method):
+            raise TypeError("RPC methods must be coroutines.")
+
+        self._rpc.add_methods((prefix, method))
+
+    def add_multi_method(self, *methods, prefix: str = None):
+        if not all(asyncio.iscoroutinefunction(m) for m in methods):
+            raise TypeError("RPC methods must be coroutines.")
+
+        for method in methods:
+            self.add_method(method, prefix=prefix)
+
+    def remove_method(self, method):
+        self._rpc.remove_method(method)
+
+    def remove_methods(self, prefix: str):
+        self._rpc.remove_methods(prefix)
+
+
+class RPCMixin:
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.rpc = RPC()
+
+        self.rpc_handlers = {}  # Uppercase cog name to method
+
+    async def _pre_login(self) -> None:
+        await self.rpc._pre_login()
+
+    def register_rpc_handler(self, method):
+        """
+        Registers a method to act as an RPC handler if the internal RPC server is active.
+
+        When calling this method through the RPC server, use the naming scheme
+        "cogname__methodname".
+
+        .. important::
+
+            All parameters to RPC handler methods must be JSON serializable objects.
+            The return value of handler methods must also be JSON serializable.
+
+        .. important::
+            RPC support is included in Red on a provisional basis. Backwards incompatible changes (up to and including removal of the RPC) may occur if deemed necessary.
+
+        Parameters
+        ----------
+        method : coroutine
+            The method to register with the internal RPC server.
+        """
+        self.rpc.add_method(method)
+
+        cog_name = method.__self__.__class__.__name__.upper()
+
+        if cog_name not in self.rpc_handlers:
+            self.rpc_handlers[cog_name] = []
+
+        self.rpc_handlers[cog_name].append(method)
+
+    def unregister_rpc_handler(self, method):
+        """
+        Deregisters an RPC method handler.
+
+        This will be called automatically for you on cog unload and will pass silently if the
+        method is not previously registered.
+
+        .. important::
+            RPC support is included in Red on a provisional basis. Backwards incompatible changes (up to and including removal of the RPC) may occur if deemed necessary.
+
+        Parameters
+        ----------
+        method : coroutine
+            The method to unregister from the internal RPC server.
+        """
+        self.rpc.remove_method(method)
+
+        name = get_name(method)
+        cog_name = name.split("__")[0]
+
+        if cog_name in self.rpc_handlers:
+            try:
+                self.rpc_handlers[cog_name].remove(method)
+            except ValueError:
+                pass