Merge pull request #771 from dmamelin/feature/stub-generator

Add IDE stubs generator service, docs, and tests

Author: craigbarratt

README.md

@@ -19,6 +19,9 @@ Pyscript also provides a kernel that interfaces with the Jupyter front-ends (eg,
 lab and VSCode). That allows you to develop and test pyscript code interactively. Plus you can interact
 with much of HASS by looking at state variables, calling services etc.
 
+Pyscript can also generate IDE stub modules by calling the `pyscript.generate_stubs` service.
+See the “IDE Helpers” section of the docs for setup details.
+
 ## Documentation
 
 Here is the [pyscript documentation](https://hacs-pyscript.readthedocs.io/en/stable).

custom_components/pyscript/__init__.py

@@ -5,6 +5,7 @@
 import json
 import logging
 import os
+import shutil
 import time
 import traceback
 from typing import Any, Callable, Dict, List, Set, Union
@@ -37,7 +38,9 @@
     FOLDER,
     LOGGER_PATH,
     REQUIREMENTS_FILE,
+    SERVICE_GENERATE_STUBS,
     SERVICE_JUPYTER_KERNEL_START,
+    SERVICE_RESPONSE_ONLY,
     UNSUB_LISTENERS,
     WATCHDOG_TASK,
 )
@@ -49,6 +52,7 @@
 from .mqtt import Mqtt
 from .requirements import install_requirements
 from .state import State, StateVal
+from .stubs.generator import StubsGenerator
 from .trigger import TrigTime
 from .webhook import Webhook
 
@@ -300,6 +304,47 @@ async def reload_scripts_handler(call: ServiceCall) -> None:
 
     hass.services.async_register(DOMAIN, SERVICE_RELOAD, reload_scripts_handler)
 
+    async def generate_stubs_service(call: ServiceCall) -> Dict[str, Any]:
+        """Generate pyscript IDE stub files."""
+
+        generator = StubsGenerator(hass)
+        generated_body = await generator.build()
+        stubs_path = os.path.join(hass.config.path(FOLDER), "modules", "stubs")
+
+        def write_stubs(path) -> dict[str, Any]:
+            res: dict[str, Any] = {}
+            try:
+                os.makedirs(path, exist_ok=True)
+
+                builtins_path = os.path.join(os.path.dirname(__file__), "stubs", "pyscript_builtins.py")
+                shutil.copy2(builtins_path, path)
+
+                gen_path = os.path.join(path, "pyscript_generated.py")
+                with open(gen_path, "w", encoding="utf-8") as f:
+                    f.write(generated_body)
+                res["status"] = "OK"
+                return res
+            except Exception as e:
+                _LOGGER.exception("Stubs generation failed: %s", e)
+                res["status"] = "Error"
+                res["exception"] = str(e)
+                res["message"] = "Check pyscript logs"
+                return res
+
+        result = await hass.async_add_executor_job(write_stubs, stubs_path)
+
+        if generator.ignored_identifiers:
+            result["ignored_identifiers"] = sorted(generator.ignored_identifiers)
+
+        if result["status"] == "OK":
+            _LOGGER.info("Pyscript stubs generated to %s", stubs_path)
+
+        return result
+
+    hass.services.async_register(
+        DOMAIN, SERVICE_GENERATE_STUBS, generate_stubs_service, supports_response=SERVICE_RESPONSE_ONLY
+    )
+
     async def jupyter_kernel_start(call: ServiceCall) -> None:
         """Handle Jupyter kernel start call."""
         _LOGGER.debug("service call to jupyter_kernel_start: %s", call.data)

custom_components/pyscript/const.py

@@ -34,6 +34,7 @@
 CONF_INSTALLED_PACKAGES = "_installed_packages"
 
 SERVICE_JUPYTER_KERNEL_START = "jupyter_kernel_start"
+SERVICE_GENERATE_STUBS = "generate_stubs"
 
 LOGGER_PATH = "custom_components.pyscript"
 

custom_components/pyscript/eval.py

@@ -999,6 +999,14 @@ async def ast_importfrom(self, arg):
                     raise ModuleNotFoundError(f"module '{imp.name}' not found")
                 self.sym_table[imp.name if imp.asname is None else imp.asname] = mod
             return
+        if arg.module == "stubs" or arg.module.startswith("stubs."):
+            for imp in arg.names:
+                if imp.asname is not None:
+                    raise ModuleNotFoundError(
+                        f"from {arg.module} import {imp.name} *as {imp.asname}* not supported for stubs"
+                    )
+            _LOGGER.debug("Skipping stubs import %s", arg.module)
+            return
         mod, error_ctx = await self.global_ctx.module_import(arg.module, arg.level)
         if error_ctx:
             self.exception_obj = error_ctx.exception_obj

custom_components/pyscript/services.yaml

@@ -105,3 +105,7 @@ jupyter_kernel_start:
       default: pyscript
       selector:
         text:
+
+generate_stubs:
+  name: Generate pyscript stubs
+  description: Build a stub files combining builtin helpers with discovered entities and services.