Experimental decorator manager

README.md

@@ -62,12 +62,14 @@ this [README](https://github.com/craigbarratt/hass-pyscript-jupyter/blob/master/
 
 ## Configuration
 
-* Go to the Integrations menu in the Home Assistant Configuration UI and add `Pyscript Python scripting` from there. Alternatively, add `pyscript:` to `<config>/configuration.yaml`; pyscript has two optional configuration parameters that allow any python package to be imported if set and to expose `hass` as a variable; both default to `false`:
+* Go to the Integrations menu in the Home Assistant Configuration UI and add `Pyscript Python scripting` from there. Alternatively, add `pyscript:` to `<config>/configuration.yaml`; pyscript has three optional configuration parameters that allow any python package to be imported if set, expose `hass` as a variable, and temporarily switch back to the legacy decorator subsystem; all three default to `false`:
     ```yaml
     pyscript:
       allow_all_imports: true
       hass_is_global: true
+      legacy_decorators: true
     ```
+   Starting with version `2.0.0`, pyscript uses the new decorator subsystem by default. If you find a problem in the new implementation, you can temporarily set `legacy_decorators: true` to switch back to the legacy one. If you do, please also file a bug report in [GitHub Issues](https://github.com/custom-components/pyscript/issues) so the new subsystem can be fixed.
 * Add files with a suffix of `.py` in the folder `<config>/pyscript`.
 * Restart HASS.
 * Whenever you change a script file, make a `reload` service call to `pyscript`.

custom_components/pyscript/__init__.py

@@ -33,6 +33,7 @@
 from .const import (
     CONF_ALLOW_ALL_IMPORTS,
     CONF_HASS_IS_GLOBAL,
+    CONF_LEGACY_DECORATORS,
     CONFIG_ENTRY,
     CONFIG_ENTRY_OLD,
     DOMAIN,
@@ -44,6 +45,7 @@
     UNSUB_LISTENERS,
     WATCHDOG_TASK,
 )
+from .decorator import DecoratorRegistry
 from .eval import AstEval
 from .event import Event
 from .function import Function
@@ -62,6 +64,7 @@
     {
         vol.Optional(CONF_ALLOW_ALL_IMPORTS, default=False): cv.boolean,
         vol.Optional(CONF_HASS_IS_GLOBAL, default=False): cv.boolean,
+        vol.Optional(CONF_LEGACY_DECORATORS, default=False): cv.boolean,
     },
     extra=vol.ALLOW_EXTRA,
 )
@@ -114,14 +117,15 @@ async def update_yaml_config(hass: HomeAssistant, config_entry: ConfigEntry) ->
     # since they affect all scripts
     #
     config_save = {
-        param: config_entry.data.get(param, False) for param in [CONF_HASS_IS_GLOBAL, CONF_ALLOW_ALL_IMPORTS]
+        param: config_entry.data.get(param, False)
+        for param in [CONF_HASS_IS_GLOBAL, CONF_ALLOW_ALL_IMPORTS, CONF_LEGACY_DECORATORS]
     }
     if DOMAIN not in hass.data:
         hass.data.setdefault(DOMAIN, {})
     if CONFIG_ENTRY_OLD in hass.data[DOMAIN]:
         old_entry = hass.data[DOMAIN][CONFIG_ENTRY_OLD]
         hass.data[DOMAIN][CONFIG_ENTRY_OLD] = config_save
-        for param in [CONF_HASS_IS_GLOBAL, CONF_ALLOW_ALL_IMPORTS]:
+        for param in [CONF_HASS_IS_GLOBAL, CONF_ALLOW_ALL_IMPORTS, CONF_LEGACY_DECORATORS]:
             if old_entry.get(param, False) != config_entry.data.get(param, False):
                 return True
     hass.data[DOMAIN][CONFIG_ENTRY_OLD] = config_save
@@ -272,6 +276,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
     Webhook.init(hass)
     State.register_functions()
     GlobalContextMgr.init()
+    DecoratorRegistry.init(hass, config_entry)
 
     pyscript_folder = hass.config.path(FOLDER)
     if not await hass.async_add_executor_job(os.path.isdir, pyscript_folder):

custom_components/pyscript/config_flow.py

@@ -9,14 +9,21 @@
 from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
 from homeassistant.core import callback
 
-from .const import CONF_ALLOW_ALL_IMPORTS, CONF_HASS_IS_GLOBAL, CONF_INSTALLED_PACKAGES, DOMAIN
+from .const import (
+    CONF_ALLOW_ALL_IMPORTS,
+    CONF_HASS_IS_GLOBAL,
+    CONF_INSTALLED_PACKAGES,
+    CONF_LEGACY_DECORATORS,
+    DOMAIN,
+)
 
-CONF_BOOL_ALL = {CONF_ALLOW_ALL_IMPORTS, CONF_HASS_IS_GLOBAL}
+CONF_BOOL_ALL = (CONF_ALLOW_ALL_IMPORTS, CONF_HASS_IS_GLOBAL, CONF_LEGACY_DECORATORS)
 
 PYSCRIPT_SCHEMA = vol.Schema(
     {
         vol.Optional(CONF_ALLOW_ALL_IMPORTS, default=False): bool,
         vol.Optional(CONF_HASS_IS_GLOBAL, default=False): bool,
+        vol.Optional(CONF_LEGACY_DECORATORS, default=False): bool,
     },
     extra=vol.ALLOW_EXTRA,
 )

custom_components/pyscript/const.py

@@ -17,6 +17,7 @@
 CONF_ALLOW_ALL_IMPORTS = "allow_all_imports"
 CONF_HASS_IS_GLOBAL = "hass_is_global"
 CONF_INSTALLED_PACKAGES = "_installed_packages"
+CONF_LEGACY_DECORATORS = "legacy_decorators"
 
 SERVICE_JUPYTER_KERNEL_START = "jupyter_kernel_start"
 SERVICE_GENERATE_STUBS = "generate_stubs"

custom_components/pyscript/decorator.py

@@ -0,0 +1,283 @@
+"""Decorator registry and manager logic for pyscript decorators."""
+
+from __future__ import annotations
+
+import ast
+import asyncio
+import logging
+import os
+from typing import Any, ClassVar
+import weakref
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import Context, HomeAssistant
+
+from .const import CONF_LEGACY_DECORATORS
+from .decorator_abc import (
+    CallHandlerDecorator,
+    CallResultHandlerDecorator,
+    Decorator,
+    DecoratorManager,
+    DecoratorManagerStatus,
+    DispatchData,
+    TriggerDecorator,
+    TriggerHandlerDecorator,
+)
+from .eval import AstEval, EvalFunc, EvalFuncVar
+from .function import Function
+from .state import State
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class DecoratorRegistry:
+    """Decorator registry."""
+
+    _decorators: dict[str, type[Decorator]]  # decorator name to class
+    hass: ClassVar[HomeAssistant]
+
+    @classmethod
+    def init(cls, hass: HomeAssistant, config_entry: ConfigEntry = None) -> None:
+        """Initialize the decorator registry."""
+        cls.hass = hass
+        cls._decorators = {}
+        disabled = False
+        if config_entry is not None and config_entry.data.get(CONF_LEGACY_DECORATORS, False):
+            disabled = True
+        elif "PYTEST_CURRENT_TEST" in os.environ and "NODM" in os.environ:
+            disabled = True
+
+        if disabled:
+            _LOGGER.warning("Using legacy decorators")
+            return
+
+        DecoratorManager.hass = hass
+
+        Function.register_ast({"task.wait_until": DecoratorRegistry.wait_until_factory})
+
+        from .decorators import DECORATORS  # pylint: disable=import-outside-toplevel
+
+        for dec_type in DECORATORS:
+            cls.register(dec_type)
+
+    @classmethod
+    def register(cls, dec_type: type[Decorator]) -> None:
+        """Register a decorator."""
+        if not dec_type.name:
+            raise TypeError(f"Decorator name is required {dec_type}")
+
+        _LOGGER.debug("Registering decorator @%s %s", dec_type.name, dec_type)
+        if dec_type.name in cls._decorators:
+            _LOGGER.warning(
+                "Overriding decorator: %s %s with %s",
+                dec_type.name,
+                cls._decorators[dec_type.name],
+                dec_type,
+            )
+        cls._decorators[dec_type.name] = dec_type
+
+    @classmethod
+    async def get_decorator_by_expr(cls, ast_ctx: AstEval, dec_expr: ast.expr) -> Decorator | None:
+        """Return decorator instance from an AST decorator expression."""
+        dec_name = None
+        has_args = False
+
+        if isinstance(dec_expr, ast.Name):  # decorator without ()
+            dec_name = dec_expr.id
+        elif isinstance(dec_expr, ast.Call) and isinstance(dec_expr.func, ast.Name):
+            dec_name = dec_expr.func.id
+            has_args = True
+
+        if know_decorator := cls._decorators.get(dec_name):
+            if has_args:
+                args = await ast_ctx.eval_elt_list(dec_expr.args)
+                kwargs = {keyw.arg: await ast_ctx.aeval(keyw.value) for keyw in dec_expr.keywords}
+            else:
+                args = []
+                kwargs = {}
+
+            decorator = know_decorator(args, kwargs)
+            return decorator
+
+        return None
+
+    @classmethod
+    async def wait_until(cls, ast_ctx: AstEval, *_arg: Any, **kwargs: Any) -> Any:
+        """Build a temporary decorator manager that waits until one of trigger decorators fires."""
+        func_args = set(kwargs.keys())
+        if len(func_args) == 0:
+            return {"trigger_type": "none"}
+
+        found_args = set()
+        dm = WaitUntilDecoratorManager(ast_ctx, **kwargs)
+
+        found_args.add("timeout")
+        found_args.add("__test_handshake__")
+
+        for dec_name, dec_class in cls._decorators.items():
+            if not issubclass(dec_class, TriggerDecorator):
+                continue
+            if dec_name not in func_args:
+                continue
+
+            dec_args = kwargs[dec_name]
+            if not isinstance(dec_args, list):
+                dec_args = [dec_args]
+            found_args.add(dec_name)
+
+            dec_kwargs = {}
+            func_args.remove(dec_name)
+            kwargs_schema_keys = dec_class.kwargs_schema.schema.keys()
+            for key in kwargs_schema_keys:
+                if key in kwargs:
+                    dec_kwargs[key] = kwargs[key]
+                    found_args.add(key)
+            dec = dec_class(dec_args, dec_kwargs)
+            dm.add(dec)
+
+        unknown_args = set(kwargs.keys()).difference(found_args)
+        if unknown_args:
+            raise ValueError(f"Unknown arguments: {unknown_args}")
+        await dm.validate()
+
+        # state_trigger sets __test_handshake__ after the initial checks.
+        # In some cases, it returns a value before __test_handshake__ is set.
+        if "state_trigger" not in kwargs:
+            if test_handshake := kwargs.get("__test_handshake__"):
+                #
+                # used for testing to avoid race conditions
+                # we use this as a handshake that we are about to
+                # listen to the queue
+                #
+                State.set(test_handshake[0], test_handshake[1])
+        await dm.start()
+
+        ret = await dm.wait_until()
+
+        return ret
+
+    @classmethod
+    def wait_until_factory(cls, ast_ctx):
+        """Return wrapper to call to astFunction with the ast context."""
+
+        async def wait_until_call(*arg, **kw):
+            return await cls.wait_until(ast_ctx, *arg, **kw)
+
+        return wait_until_call
+
+
+class WaitUntilDecoratorManager(DecoratorManager):
+    """Decorator manager for task.wait_until."""
+
+    def __init__(self, ast_ctx: AstEval, **kwargs: dict[str, Any]) -> None:
+        """Initialize the task.wait_until decorator manager."""
+        super().__init__(ast_ctx, ast_ctx.name)
+        self.kwargs = kwargs
+        self._future: asyncio.Future[DispatchData] = self.hass.loop.create_future()
+        self.timeout_decorator = None
+        if timeout := kwargs.get("timeout"):
+            to_dec = DecoratorRegistry._decorators.get("time_trigger")
+            self.timeout_decorator = to_dec([f"once(now + {timeout}s)"], {})
+            self.add(self.timeout_decorator)
+
+    async def dispatch(self, data: DispatchData) -> None:
+        """Resolve the waiting future on the first incoming dispatch."""
+        _LOGGER.debug("task.wait_until dispatch: %s", data)
+        if self._future.done():
+            _LOGGER.debug("task.wait_until future already completed: %s", self._future.exception())
+            # ignore another calls
+            return
+        await self.stop()
+        self._future.set_result(data)
+
+    async def handle_exception(self, exc: Exception) -> None:
+        """Propagate an evaluation exception to the waiting caller."""
+        if self._future.done():
+            _LOGGER.debug("task.wait_until future already completed: %s", self._future.exception())
+            return
+        await self.stop()
+        self._future.set_exception(exc)
+
+    async def wait_until(self) -> dict[str, Any]:
+        """Wait for dispatch and normalize the return payload."""
+        data = await self._future
+        if data.trigger == self.timeout_decorator:
+            ret = {"trigger_type": "timeout"}
+        else:
+            ret = data.func_args
+        _LOGGER.debug("task.wait_until finish: %s", ret)
+        return ret
+
+
+class FunctionDecoratorManager(DecoratorManager):
+    """Maintain and validate a set of decorators applied to a function."""
+
+    def __init__(self, ast_ctx: AstEval, eval_func_var: EvalFuncVar) -> None:
+        """Initialize the function decorator manager."""
+        super().__init__(ast_ctx, f"{ast_ctx.get_global_ctx_name()}.{eval_func_var.get_name()}")
+        self.eval_func: EvalFunc = eval_func_var.func
+
+        self.logger = self.eval_func.logger
+
+        def on_func_var_deleted():
+            if self.status is DecoratorManagerStatus.RUNNING:
+                self.hass.async_create_task(self.stop())
+
+        weakref.finalize(eval_func_var, on_func_var_deleted)
+
+    async def _call(self, data: DispatchData) -> None:
+        handlers = self.get_decorators(CallHandlerDecorator)
+        result_handlers = self.get_decorators(CallResultHandlerDecorator)
+
+        for handler_dec in handlers:
+            if await handler_dec.handle_call(data) is False:
+                self.logger.debug("Calling canceled by %s", handler_dec)
+                # notify handlers with "None"
+                for result_handler_dec in result_handlers:
+                    await result_handler_dec.handle_call_result(data, None)
+                return
+        # Fire an event indicating that pyscript is running
+        # Note: the event must have an entity_id for logbook to work correctly.
+        ev_name = self.name.replace(".", "_")
+        ev_entity_id = f"pyscript.{ev_name}"
+
+        event_data = {"name": ev_name, "entity_id": ev_entity_id, "func_args": data.func_args}
+        self.hass.bus.async_fire("pyscript_running", event_data, context=data.hass_context)
+        # Store HASS Context for this Task
+        Function.store_hass_context(data.hass_context)
+
+        result = await data.call_ast_ctx.call_func(self.eval_func, None, **data.func_args)
+        for result_handler_dec in result_handlers:
+            await result_handler_dec.handle_call_result(data, result)
+
+    async def dispatch(self, data: DispatchData) -> None:
+        """Handle a trigger dispatch: run guards, create a context, and invoke the function."""
+        _LOGGER.debug("Dispatching for %s: %s", self.name, data)
+
+        decorators = self.get_decorators(TriggerHandlerDecorator)
+        for dec in decorators:
+            if await dec.handle_dispatch(data) is False:
+                self.logger.debug("Trigger not active due to %s", dec)
+                return
+
+        action_ast_ctx = AstEval(
+            f"{self.eval_func.global_ctx_name}.{self.eval_func.name}", self.eval_func.global_ctx
+        )
+        Function.install_ast_funcs(action_ast_ctx)
+        data.call_ast_ctx = action_ast_ctx
+
+        # Create new HASS Context with incoming as parent
+        if "context" in data.func_args and isinstance(data.func_args["context"], Context):
+            data.hass_context = Context(parent_id=data.func_args["context"].id)
+        else:
+            data.hass_context = Context()
+
+        self.logger.debug(
+            "trigger %s got %s trigger, running action (kwargs = %s)",
+            self.name,
+            data.trigger,
+            data.func_args,
+        )
+
+        task = Function.create_task(self._call(data), ast_ctx=action_ast_ctx)
+        Function.task_done_callback_ctx(task, action_ast_ctx)