From e3755685356904842f31ba680b827b3dd024bef2 Mon Sep 17 00:00:00 2001 From: Easton Potokar Date: Sat, 1 Aug 2026 15:12:04 -0400 Subject: [PATCH] Clean up some type generation --- custom_components/pyscript/stubs/generator.py | 37 ++++-- .../pyscript/stubs/pyscript_builtins.py | 9 +- tests/test_stubs.py | 107 ++++++++++++++++-- 3 files changed, 130 insertions(+), 23 deletions(-) diff --git a/custom_components/pyscript/stubs/generator.py b/custom_components/pyscript/stubs/generator.py index 8f5aa36..7a00eb0 100644 --- a/custom_components/pyscript/stubs/generator.py +++ b/custom_components/pyscript/stubs/generator.py @@ -11,7 +11,6 @@ from custom_components.pyscript.stubs.pyscript_builtins import StateVal from homeassistant.core import HomeAssistant, split_entity_id -from homeassistant.helpers import entity_registry as er from homeassistant.helpers.service import async_get_all_descriptions _LOGGER = logging.getLogger(__name__) @@ -70,10 +69,11 @@ async def build(self) -> str: } for module, imports in base_imports.items(): + level = 1 if module == "pyscript_builtins" else 0 module_body.append( ast.ImportFrom( module=module, - level=0, + level=level, names=[ast.alias(name=imp, asname=None) for imp in imports], ) ) @@ -163,13 +163,11 @@ def _get_entity_class_name(self, domain_id: str) -> str: return f"_{domain_id}{_STATE_CLASS_SUFFIX}" async def _build_entity_classes(self): - for entity in er.async_get(self._hass).entities.values(): - if entity.disabled: - continue - - domain_id, entity_id = split_entity_id(entity.entity_id) + for state_obj in sorted(self._hass.states.async_all(), key=lambda s: s.entity_id): + full_entity_id = state_obj.entity_id + domain_id, entity_id = split_entity_id(full_entity_id) - if not self._is_identifier(entity_id, entity.entity_id): + if not self._is_identifier(entity_id, full_entity_id): continue self._collect_entity_atts(domain_id, entity_id) @@ -235,6 +233,7 @@ async def _create_service_function( kwonlyargs: list[ast.arg] = [] kw_defaults: list[ast.expr] = [] decorator_list: list[ast.expr] = [] + defaults: list[ast.expr] = [] has_target = "target" in payload @@ -258,6 +257,9 @@ async def _create_service_function( if def_type == "entity" and len(field_nodes) == 1: # simple calling with 1 arg service args.append(ast.arg(arg=field_nodes[0].name, annotation=field_nodes[0].annotation)) + # Preserve the default for the single positional argument. + if field_nodes[0].default is not None: + defaults.append(field_nodes[0].default) else: for field in field_nodes: kwonlyargs.append(ast.arg(arg=field.name, annotation=field.annotation)) @@ -287,7 +289,7 @@ async def _create_service_function( kwonlyargs=kwonlyargs, kw_defaults=kw_defaults, kwarg=None, - defaults=[], + defaults=defaults, ), body=body, decorator_list=decorator_list, @@ -330,6 +332,13 @@ def _describe_service_field( if default_value is not None and isinstance(default_value, (int, float, str, bool)): default_expr = ast.Constant(value=default_value) + # Widen annotation when the default's type conflicts with the selector type. + if default_expr is not None and annotation is not None: + ann_str = ast.unparse(annotation) + default_type = type(default_value).__name__ + if default_type not in ann_str and default_type in ("int", "float", "str", "bool"): + annotation = ast.BinOp(left=annotation, op=ast.BitOr(), right=self._name(default_type)) + if not is_required: if default_expr is None: if annotation is not None: @@ -361,8 +370,14 @@ def _selector_annotation(self, selector: dict[str, Any] | None) -> ast.expr | No if selector_id == "number": if selector_value == "any": return self._name("float") - if isinstance(selector_value, dict) and selector_value.get("mode") == "box": - return self._name("float") + if isinstance(selector_value, dict): + if selector_value.get("mode") == "box": + return self._name("float") + # Use float when step or min/max are fractional. + for key in ("step", "min", "max"): + val = selector_value.get(key) + if isinstance(val, float) and val != int(val): + return self._name("float") return self._name("int") if selector_id == "select": options = [] diff --git a/custom_components/pyscript/stubs/pyscript_builtins.py b/custom_components/pyscript/stubs/pyscript_builtins.py index f7386d5..ee77592 100644 --- a/custom_components/pyscript/stubs/pyscript_builtins.py +++ b/custom_components/pyscript/stubs/pyscript_builtins.py @@ -12,11 +12,12 @@ from datetime import datetime from typing import Any, Literal -from homeassistant.components.webhook import SUPPORTED_METHODS from homeassistant.core import HomeAssistant hass: HomeAssistant +WebhookMethod = Literal["GET", "HEAD", "POST", "PUT"] + def service( *service_name: str, supports_response: Literal["none", "only", "optional"] = "none" @@ -129,7 +130,7 @@ def webhook_trigger( webhook_id: str, str_expr: str | None = None, local_only: bool = True, - methods: set[SUPPORTED_METHODS] | list[SUPPORTED_METHODS] = {"POST", "PUT"}, + methods: set[WebhookMethod] | list[WebhookMethod] = {"POST", "PUT"}, kwargs: dict | None = None, ) -> Callable[..., Any]: """Trigger when a request is made to a webhook endpoint. @@ -150,7 +151,7 @@ def webhook_handler( webhook_id: str, str_expr: str | None = None, local_only: bool = True, - methods: set[SUPPORTED_METHODS] | list[SUPPORTED_METHODS] = {"POST", "PUT"}, + methods: set[WebhookMethod] | list[WebhookMethod] = {"POST", "PUT"}, timeout: int | float = 10.0, kwargs: dict | None = None, ) -> Callable[..., Any]: @@ -491,7 +492,7 @@ def wait_until( mqtt_trigger_encoding: str | None = None, webhook_trigger: str | list[str] | None = None, webhook_local_only: bool = True, - webhook_methods: list[SUPPORTED_METHODS] = ("POST", "PUT"), + webhook_methods: list[WebhookMethod] = ["POST", "PUT"], timeout: int | float | None = None, state_check_now: bool = True, state_hold: int | float | None = None, diff --git a/tests/test_stubs.py b/tests/test_stubs.py index d1b3650..fdaea93 100644 --- a/tests/test_stubs.py +++ b/tests/test_stubs.py @@ -4,7 +4,6 @@ from datetime import datetime as dt from pathlib import Path -from types import SimpleNamespace from typing import Any import pytest @@ -40,12 +39,7 @@ def ready(): }, ) - dummy_registry = SimpleNamespace( - entities={ - "light.lamp": SimpleNamespace(entity_id="light.lamp", disabled=False), - } - ) - monkeypatch.setattr("custom_components.pyscript.stubs.generator.er.async_get", lambda _: dummy_registry) + hass.states.async_set("remote.tv", "off", {}) async def fake_service_descriptions(_hass: HomeAssistant) -> dict[str, dict[str, dict[str, Any]]]: return { @@ -68,7 +62,44 @@ async def fake_service_descriptions(_hass: HomeAssistant) -> dict[str, dict[str, }, "response": {"optional": True}, } - } + }, + "remote": { + "send_command": { + "description": "Send a command.", + "fields": { + "delay_secs": { + "required": False, + "default": 0.4, + "selector": {"number": {"step": 0.1, "min": 0.0, "max": 5.0}}, + "description": "Delay between commands.", + }, + }, + }, + "turn_on": { + "description": "Turn on remote.", + "target": {"entity": {"domain": "remote"}}, + "fields": { + "activity": { + "required": False, + "selector": {"text": None}, + "description": "Activity to start.", + }, + }, + }, + }, + "voice_satellite": { + "show": { + "description": "Show something.", + "fields": { + "pipeline": { + "required": False, + "default": 1, + "selector": {"text": None}, + "description": "Pipeline name or slot.", + }, + }, + }, + }, } monkeypatch.setattr( @@ -115,6 +146,13 @@ async def fake_service_descriptions(_hass: HomeAssistant) -> dict[str, dict[str, assert "'fast'" in generated_content assert "-> dict[str, Any]" in generated_content + # Fractional number selector produces float annotation. + assert "delay_secs: float" in generated_content + # Type mismatch between selector (str) and default (int) widens annotation. + assert "str | int" in generated_content + # Entity service with single optional arg preserves the default. + assert "turn_on(self, activity: str | None=None)" in generated_content + original_builtins = ( Path(__file__).resolve().parent.parent / "custom_components" @@ -156,3 +194,56 @@ def stub_import_ready(): assert hass.services.has_service(DOMAIN, "stub_import_ready") assert "ModuleNotFoundError" not in caplog.text + + +@pytest.mark.asyncio +async def test_stubs_include_state_only_entities(hass, caplog, monkeypatch): + """Entities only in the state machine (not in entity registry) must appear in stubs.""" + + await setup_script( + hass, + notify_q=None, + now=dt(2024, 3, 3, 0, 0, 0), + source=""" +@service +def ready2(): + pass +""", + script_name="/stub_state_only.py", + ) + + # Entity in state machine but NOT in the registry (template entity / zone.home pattern). + hass.states.async_set( + "binary_sensor.my_sensor", + "off", + {"device_class": "running", "friendly_name": "Van Car Running"}, + ) + hass.states.async_set( + "zone.home", + "0", + {"latitude": 40.0, "longitude": -74.0, "radius": 100, "persons": ["person.me"]}, + ) + + # Empty registry - forces all entities to come from state machine only. + + async def fake_service_descriptions(_hass: HomeAssistant) -> dict[str, dict[str, dict[str, Any]]]: + return {} + + monkeypatch.setattr( + "custom_components.pyscript.stubs.generator.async_get_all_descriptions", fake_service_descriptions + ) + + stubs_dir = Path(hass.config.path(FOLDER)) / "modules" / "stubs" + generated_target = stubs_dir / "pyscript_generated.py" + stubs_dir.mkdir(parents=True, exist_ok=True) + + await hass.services.async_call(DOMAIN, SERVICE_GENERATE_STUBS, {}, blocking=True, return_response=True) + + content = generated_target.read_text(encoding="utf-8") + assert "my_sensor: _binary_sensor_state" in content + assert "home: _zone_state" in content + + # Cleanup + for child in stubs_dir.iterdir(): + child.unlink() + stubs_dir.rmdir()