revamped autoloading and importing as discussed in #26 and fixes #28

Author: craigbarratt

custom_components/pyscript/__init__.py

@@ -64,7 +64,7 @@ def check_isdir(path):
 
     State.set_pyscript_config(config.get(DOMAIN, {}))
 
-    await load_scripts(hass)
+    await load_scripts(hass, config)
 
     async def reload_scripts_handler(call):
         """Handle reload service calls."""
@@ -80,7 +80,6 @@ async def reload_scripts_handler(call):
 
         config = await async_process_component_config(hass, conf, integration)
 
-        # GlobalContext.global_sym_table_add("pyscript.config", config.get(DOMAIN, {}))
         State.set_pyscript_config(config.get(DOMAIN, {}))
 
         ctx_delete = {}
@@ -89,18 +88,22 @@ async def reload_scripts_handler(call):
             if idx < 0 or global_ctx_name[0:idx] not in {"file", "apps", "modules"}:
                 continue
             global_ctx.stop()
-            global_ctx.set_auto_start(False)
             ctx_delete[global_ctx_name] = global_ctx
         for global_ctx_name, global_ctx in ctx_delete.items():
             await GlobalContextMgr.delete(global_ctx_name)
 
-        await load_scripts(hass)
+        await load_scripts(hass, config)
 
         for global_ctx_name, global_ctx in GlobalContextMgr.items():
             idx = global_ctx_name.find(".")
             if idx < 0 or global_ctx_name[0:idx] not in {"file", "apps"}:
                 continue
             global_ctx.set_auto_start(True)
+
+        for global_ctx_name, global_ctx in GlobalContextMgr.items():
+            idx = global_ctx_name.find(".")
+            if idx < 0 or global_ctx_name[0:idx] not in {"file", "apps"}:
+                continue
             global_ctx.start()
 
     hass.services.async_register(DOMAIN, SERVICE_RELOAD, reload_scripts_handler)
@@ -150,10 +153,15 @@ async def start_triggers(event):
         _LOGGER.debug("adding state changed listener and starting triggers")
         hass.bus.async_listen(EVENT_STATE_CHANGED, state_changed)
         for global_ctx_name, global_ctx in GlobalContextMgr.items():
-            if not global_ctx_name.startswith("file."):
+            idx = global_ctx_name.find(".")
+            if idx < 0 or global_ctx_name[0:idx] not in {"file", "apps"}:
                 continue
-            global_ctx.start()
             global_ctx.set_auto_start(True)
+        for global_ctx_name, global_ctx in GlobalContextMgr.items():
+            idx = global_ctx_name.find(".")
+            if idx < 0 or global_ctx_name[0:idx] not in {"file", "apps"}:
+                continue
+            global_ctx.start()
 
     async def stop_triggers(event):
         _LOGGER.debug("stopping triggers")
@@ -169,25 +177,46 @@ async def stop_triggers(event):
 
 
 @bind_hass
-async def load_scripts(hass):
+async def load_scripts(hass, config):
     """Load all python scripts in FOLDER."""
 
-    load_paths = [hass.config.path(FOLDER) + "/apps", hass.config.path(FOLDER)]
+    pyscript_dir = hass.config.path(FOLDER)
 
-    _LOGGER.debug("load_scripts: load_paths = %s", load_paths)
-
-    def glob_files(load_paths, match):
+    def glob_files(load_paths, config):
         source_files = []
-        for path in load_paths:
-            source_files += sorted(glob.glob(os.path.join(path, match)))
+        apps_config = config.get(DOMAIN, {}).get("apps", None)
+        for path, match, check_config in load_paths:
+            for this_path in sorted(glob.glob(os.path.join(pyscript_dir, path, match))):
+                rel_import_path = None
+                elts = this_path.split("/")
+                if match.find("/") < 0:
+                    # last entry without the .py
+                    mod_name = elts[-1][0:-3]
+                else:
+                    # 2nd last entry
+                    mod_name = elts[-2]
+                    rel_import_path = f"{path}/mod_name"
+                if path == "":
+                    global_ctx_name = f"file.{mod_name}"
+                else:
+                    global_ctx_name = f"{path}.{mod_name}"
+                if check_config:
+                    _LOGGER.debug("load_scripts: checking %s in %s", mod_name, apps_config)
+                    if not isinstance(apps_config, dict) or mod_name not in apps_config:
+                        _LOGGER.debug("load_scripts: skipping %s because config not present", this_path)
+                        continue
+                source_files.append([global_ctx_name, this_path, rel_import_path])
         return source_files
 
-    source_files = await hass.async_add_executor_job(glob_files, load_paths, "*.py")
-
-    for source_file in sorted(source_files):
-        name = os.path.splitext(os.path.basename(source_file))[0]
-
-        global_ctx = GlobalContext(f"file.{name}", global_sym_table={}, manager=GlobalContextMgr)
-        global_ctx.set_auto_start(False)
+    load_paths = [
+        ["apps", "*.py", True],
+        ["apps", "*/__init__.py", True],
+        ["", "*.py", False],
+    ]
 
+    source_files = await hass.async_add_executor_job(glob_files, load_paths, config)
+    for global_ctx_name, source_file, rel_import_path in source_files:
+        global_ctx = GlobalContext(
+            global_ctx_name, global_sym_table={}, manager=GlobalContextMgr, rel_import_path=rel_import_path
+        )
         await GlobalContextMgr.load_file(source_file, global_ctx)

custom_components/pyscript/eval.py

@@ -479,8 +479,8 @@ def trigger_stop(self):
     async def eval_decorators(self, ast_ctx):
         """Evaluate the function decorators arguments."""
         self.decorators = []
-        ast_ctx.code_str = self.code_str
-        ast_ctx.code_list = self.code_list
+        code_str, code_list = ast_ctx.code_str, ast_ctx.code_list
+        ast_ctx.code_str, ast_ctx.code_list = self.code_str, self.code_list
         for dec in self.func_def.decorator_list:
             if isinstance(dec, ast.Call) and isinstance(dec.func, ast.Name):
                 args = []
@@ -496,6 +496,7 @@ async def eval_decorators(self, ast_ctx):
                 self.decorators.append([dec.id, None, None])
             else:
                 _LOGGER.error("function %s has unexpected decorator type %s", self.name, dec)
+        ast_ctx.code_str, ast_ctx.code_list = code_str, code_list
 
     async def resolve_nonlocals(self, ast_ctx):
         """Tag local variables and resolve nonlocals."""
@@ -624,8 +625,8 @@ async def call(self, ast_ctx, args=None, kwargs=None):
                 sym_table[name] = EvalLocalVar(name)
         ast_ctx.sym_table_stack.append(ast_ctx.sym_table)
         ast_ctx.sym_table = sym_table
-        ast_ctx.code_str = self.code_str
-        ast_ctx.code_list = self.code_list
+        code_str, code_list = ast_ctx.code_str, ast_ctx.code_list
+        ast_ctx.code_str, ast_ctx.code_list = self.code_str, self.code_list
         self.exception = None
         self.exception_obj = None
         self.exception_long = None
@@ -642,6 +643,7 @@ async def call(self, ast_ctx, args=None, kwargs=None):
                 break
         ast_ctx.sym_table = ast_ctx.sym_table_stack.pop()
         ast_ctx.curr_func = prev_func
+        ast_ctx.code_str, ast_ctx.code_list = code_str, code_list
         return val
 
 
@@ -732,7 +734,12 @@ async def ast_module(self, arg):
     async def ast_import(self, arg):
         """Execute import."""
         for imp in arg.names:
-            mod = await self.global_ctx.module_import(imp.name)
+            mod, error_ctx = await self.global_ctx.module_import(imp.name, 0)
+            if error_ctx:
+                self.exception_obj = error_ctx.exception_obj
+                self.exception = error_ctx.exception
+                self.exception_long = error_ctx.exception_long
+                raise self.exception_obj
             if not mod:
                 if not self.allow_all_imports and imp.name not in ALLOWED_IMPORTS:
                     raise ModuleNotFoundError(f"import of {imp.name} not allowed")
@@ -744,7 +751,26 @@ async def ast_import(self, arg):
 
     async def ast_importfrom(self, arg):
         """Execute from X import Y."""
-        mod = await self.global_ctx.module_import(arg.module)
+        if arg.module is None:
+            # handle: "from . import xyz"
+            for imp in arg.names:
+                mod, error_ctx = await self.global_ctx.module_import(imp.name, arg.level)
+                if error_ctx:
+                    self.exception_obj = error_ctx.exception_obj
+                    self.exception = error_ctx.exception
+                    self.exception_long = error_ctx.exception_long
+                    raise self.exception_obj
+                if not mod:
+                    raise ModuleNotFoundError(f"module '{imp.name}' not found")
+                self.sym_table[imp.name if imp.asname is None else imp.asname] = mod
+            return
+        else:
+            mod, error_ctx = await self.global_ctx.module_import(arg.module, arg.level)
+            if error_ctx:
+                self.exception_obj = error_ctx.exception_obj
+                self.exception = error_ctx.exception
+                self.exception_long = error_ctx.exception_long
+                raise self.exception_obj
         if not mod:
             if not self.allow_all_imports and arg.module not in ALLOWED_IMPORTS:
                 raise ModuleNotFoundError(f"import from {arg.module} not allowed")

custom_components/pyscript/global_ctx.py

@@ -15,7 +15,7 @@
 class GlobalContext:
     """Define class for global variables and trigger context."""
 
-    def __init__(self, name, global_sym_table=None, manager=None):
+    def __init__(self, name, global_sym_table=None, manager=None, rel_import_path=None):
         """Initialize GlobalContext."""
         self.name = name
         self.global_sym_table = global_sym_table if global_sym_table else {}
@@ -25,6 +25,7 @@ def __init__(self, name, global_sym_table=None, manager=None):
         self.manager = manager
         self.auto_start = False
         self.module = None
+        self.rel_import_path = rel_import_path
 
     def trigger_register(self, func):
         """Register a trigger function; return True if start now."""
@@ -50,11 +51,12 @@ def start(self):
         self.triggers_delay_start = set()
 
     def stop(self):
-        """Stop all triggers."""
+        """Stop all triggers and auto_start."""
         for func in self.triggers:
             func.trigger_stop()
         self.triggers = set()
         self.triggers_delay_start = set()
+        self.set_auto_start(False)
 
     def get_name(self):
         """Return the global context name."""
@@ -68,27 +70,97 @@ def get_trig_info(self, name, trig_args):
         """Return a new trigger info instance with the given args."""
         return TrigInfo(name, trig_args, self)
 
-    async def module_import(self, module_name):
-        """Import a module from the pyscript/modules folder."""
-        mod_ctx_name = f"module.{module_name}"
-        mod_ctx = self.manager.get(mod_ctx_name)
-        if not mod_ctx:
-
-            def check_isfile(filepath):
-                return os.path.isfile(filepath)
-
-            filepath = Function.hass.config.path(FOLDER) + f"/modules/{module_name}.py"
-            if await Function.hass.async_add_executor_job(check_isfile, filepath):
-                mod = ModuleType(module_name)
-                global_ctx = GlobalContext(mod_ctx_name, global_sym_table=mod.__dict__, manager=self.manager)
-                global_ctx.set_auto_start(True)
-                await self.manager.load_file(filepath, global_ctx)
-                global_ctx.module = mod
-                mod_ctx = self.manager.get(mod_ctx_name)
-
-        if mod_ctx and mod_ctx.module:
-            return mod_ctx.module
-        return None
+    async def module_import(self, module_name, import_level):
+        """Import a pyscript module from the pyscript/modules or apps folder."""
+
+        pyscript_dir = Function.hass.config.path(FOLDER)
+        module_path = module_name.replace(".", "/")
+        filepaths = []
+
+        def find_first_file(filepaths):
+            for ctx_name, path, rel_path in filepaths:
+                abs_path = os.path.join(pyscript_dir, path)
+                if os.path.isfile(abs_path):
+                    return [ctx_name, abs_path, rel_path]
+            return None
+
+        #
+        # first build a list of potential import files
+        #
+        if import_level > 0:
+            if self.rel_import_path is None:
+                raise ImportError("attempted relative import with no known parent package")
+            path = self.rel_import_path
+            ctx_name = self.name
+            for _ in range(import_level - 1):
+                path = os.path.basename(path)
+                idx = ctx_name.rfind(".")
+                if path.find("/") < 0 or idx < 0:
+                    raise ImportError("attempted relative import above parent package")
+                ctx_name = ctx_name[0:idx]
+            idx = ctx_name.rfind(".")
+            if idx >= 0:
+                ctx_name = f"{ctx_name[0:idx]}.{module_name}"
+            filepaths.append([ctx_name, f"{path}/{module_path}.py", path])
+            path += f"/{module_path}"
+            filepaths.append([f"{ctx_name}.__init__", f"{path}/__init__.py", path])
+            module_name = ctx_name[ctx_name.find(".") + 1 :]
+
+        else:
+            if self.rel_import_path is not None and self.rel_import_path.startswith("apps/"):
+                ctx_name = f"apps.{module_name}"
+                filepaths.append([ctx_name, f"apps/{module_path}.py", f"apps/{module_path}"])
+                filepaths.append(
+                    [f"{ctx_name}.__init__", f"apps/{module_path}/__init__.py", f"apps/{module_path}"]
+                )
+
+            ctx_name = f"modules.{module_name}"
+            filepaths.append([ctx_name, f"modules/{module_path}.py", None])
+            filepaths.append(
+                [f"{ctx_name}.__init__", f"modules/{module_path}/__init__.py", f"modules/{module_path}"]
+            )
+
+        #
+        # now see if we have loaded it already
+        #
+        for ctx_name, _, _ in filepaths:
+            mod_ctx = self.manager.get(ctx_name)
+            if mod_ctx and mod_ctx.module:
+                _LOGGER.debug(
+                    "module_import: returning existing module %s, ctx = %s, mod = %s",
+                    module_name,
+                    ctx_name,
+                    mod_ctx.module,
+                )
+                return [mod_ctx.module, None]
+
+        #
+        # not loaded already, so try to find and import it
+        #
+        _LOGGER.debug("module_import: for module %s, searching %s", module_name, filepaths)
+        file_info = await Function.hass.async_add_executor_job(find_first_file, filepaths)
+        if not file_info:
+            return [None, None]
+
+        [ctx_name, filepath, rel_import_path] = file_info
+
+        mod = ModuleType(module_name)
+        global_ctx = GlobalContext(
+            ctx_name, global_sym_table=mod.__dict__, manager=self.manager, rel_import_path=rel_import_path
+        )
+        global_ctx.set_auto_start(True)
+        error_ctx = await self.manager.load_file(filepath, global_ctx)
+        if error_ctx:
+            _LOGGER.debug(
+                "module_import: failed to load module %s, ctx = %s, path = %s",
+                module_name,
+                ctx_name,
+                filepath,
+            )
+            return [None, error_ctx]
+        global_ctx.module = mod
+        _LOGGER.debug("module_import: imported %s, ctx = %s, mod = %s", module_name, ctx_name, mod)
+        return [mod, None]
 
 
 class GlobalContextMgr:
@@ -185,7 +257,7 @@ def new_name(cls, root):
 
     @classmethod
     async def load_file(cls, filepath, global_ctx):
-        """Load, parse and run the given script file."""
+        """Load, parse and run the given script file; returns error ast_ctx on error, or None if ok."""
 
         def read_file(path):
             try:
@@ -201,7 +273,7 @@ def read_file(path):
         source = await Function.hass.async_add_executor_job(read_file, filepath)
 
         if source is None:
-            return
+            return False
 
         ast_ctx = AstEval(global_ctx.get_name(), global_ctx)
         Function.install_ast_funcs(ast_ctx)
@@ -210,11 +282,12 @@ def read_file(path):
             exc = ast_ctx.get_exception_long()
             ast_ctx.get_logger().error(exc)
             global_ctx.stop()
-            return
+            return ast_ctx
         await ast_ctx.eval()
         exc = ast_ctx.get_exception_long()
         if exc is not None:
             ast_ctx.get_logger().error(exc)
             global_ctx.stop()
-            return
+            return ast_ctx
         cls.set(global_ctx.get_name(), global_ctx)
+        return None