Merge pull request #863 from contagon/strict-types

Tweaks in type stubs generation

Author: craigbarratt

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 = []

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,

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()