diff --git a/custom_components/pyscript/stubs/generator.py b/custom_components/pyscript/stubs/generator.py index 16dd069..e37aa73 100644 --- a/custom_components/pyscript/stubs/generator.py +++ b/custom_components/pyscript/stubs/generator.py @@ -3,6 +3,7 @@ from __future__ import annotations import ast +import copy from dataclasses import dataclass from enum import IntEnum, IntFlag, StrEnum import keyword @@ -184,14 +185,21 @@ async def _build_entity_classes(self): self._get_or_create_class(self._get_entity_class_name(domain_id), _STATE_CLASS) async def _build_services(self): - def process_fields(fields: dict[str, Any]) -> list[_ServiceField]: + def process_fields(fields: dict[str, Any], seen: set[str]) -> list[_ServiceField]: result: list[_ServiceField] = [] for field_name, field in (fields.get("fields") or {}).items(): - if field_name in ("additional_fields", "advanced_fields"): - result.extend(process_fields(field)) + if self._is_section(field): + # a section groups fields in the UI only, service data stays flat + result.extend(process_fields(field, seen)) + continue + if field_name in seen: + _LOGGER.debug( + "Skipping duplicate field %s in service %s.%s", field_name, domain_id, service_id + ) continue definition = self._describe_service_field(service_id, field_name, field) if definition is not None: + seen.add(field_name) result.append(definition) return result @@ -204,8 +212,9 @@ def process_fields(fields: dict[str, Any]) -> list[_ServiceField]: _LOGGER.debug("Building service %s.%s payload: %s", domain_id, service_id, payload) + # services with a target get an implicit entity_id argument field_nodes = sorted( - process_fields(payload), + process_fields(payload, {"entity_id"} if "target" in payload else set()), key=lambda x: not x.required, ) @@ -334,9 +343,9 @@ def _describe_service_field( # 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) + ann_names = {node.id for node in ast.walk(annotation) if isinstance(node, ast.Name)} default_type = type(default_value).__name__ - if default_type not in ann_str and default_type in ("int", "float", "str", "bool"): + if default_type not in ann_names 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: @@ -365,35 +374,60 @@ def _selector_annotation(self, selector: dict[str, Any] | None) -> ast.expr | No if not selector: return None for selector_id, selector_value in selector.items(): - if selector_type := SELECTOR_SIMPLE_TYPES.get(selector_id): - return self._name(selector_type) - if selector_id == "number": - if selector_value == "any": - 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 = [] - if isinstance(selector_value, dict): - options = selector_value.get("options") or [] - literals = [ast.Constant(value="")] - for option in options: - value = option.get("value") if isinstance(option, dict) else option - literals.append(ast.Constant(value=value)) - return ast.Subscript( - value=self._name("Literal"), - slice=ast.Tuple(elts=literals), + annotation = self._selector_base_annotation(selector_id, selector_value) + if annotation is None: + continue + if isinstance(selector_value, dict) and selector_value.get("multiple") is True: + # Home Assistant also accepts a single value for "multiple" selectors (cv.ensure_list) + return ast.BinOp( + left=annotation, + op=ast.BitOr(), + right=ast.Subscript(value=self._name("list"), slice=copy.deepcopy(annotation)), ) + return annotation _LOGGER.debug("Selector annotation unknown %s", selector) return None + def _selector_base_annotation(self, selector_id: str, selector_value: Any) -> ast.expr | None: + if selector_type := SELECTOR_SIMPLE_TYPES.get(selector_id): + return self._name(selector_type) + if selector_id == "number": + if selector_value == "any": + 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 = [] + if isinstance(selector_value, dict): + options = selector_value.get("options") or [] + literals = [ast.Constant(value="")] + for option in options: + value = option.get("value") if isinstance(option, dict) else option + literals.append(ast.Constant(value=value)) + return ast.Subscript( + value=self._name("Literal"), + slice=ast.Tuple(elts=literals), + ) + return None + + @staticmethod + def _is_section(field: Any) -> bool: + """Return True if a service field is a UI section grouping other fields.""" + if not isinstance(field, dict) or "selector" in field: + return False + sub_fields = field.get("fields") + # an empty section is still a section, it just contributes no fields + return isinstance(sub_fields, dict) and all( + isinstance(sub_field, dict) for sub_field in sub_fields.values() + ) + def _is_identifier(self, value: str, fqn: str) -> bool: valid = value.isidentifier() and not keyword.iskeyword(value) if not valid: diff --git a/tests/test_stubs.py b/tests/test_stubs.py index 10dd3d8..bfc5e97 100644 --- a/tests/test_stubs.py +++ b/tests/test_stubs.py @@ -313,3 +313,132 @@ async def fake_service_descriptions(_hass: HomeAssistant) -> dict[str, dict[str, for child in stubs_dir.iterdir(): child.unlink() stubs_dir.rmdir() + + +@pytest.mark.asyncio +async def test_stubs_flatten_named_sections_and_multiple(pyscript, monkeypatch): + """Sections of any name must be flattened, and "multiple" selectors accept a value or a list.""" + + hass = pyscript.hass + + async def fake_service_descriptions(_hass: HomeAssistant) -> dict[str, dict[str, dict[str, Any]]]: + return { + "mailer": { + "send": { + "description": "Send a message.", + "fields": { + "recipients": { + "required": True, + "selector": {"text": {"multiple": True}}, + "description": "Recipients.", + }, + "sender": { + "required": False, + "default": "admin", + "selector": {"text": {"multiple": True}}, + "description": "Sender.", + }, + "subject": { + "required": False, + "selector": {"text": None}, + "description": "Subject.", + }, + "legacy": { + "required": False, + "fields": "not a section", + "description": "Legacy field with a fields key.", + }, + "delivery_options": { + "collapsed": True, + "fields": { + "channels": { + "required": False, + "selector": {"select": {"options": ["email", "sms"], "multiple": True}}, + "description": "Channels.", + }, + "camera": { + "required": False, + "selector": {"entity": {"domain": "camera"}}, + "description": "Camera.", + }, + "subject": { + "required": False, + "selector": {"number": None}, + "description": "Duplicate of a top-level field.", + }, + }, + }, + "advanced": { + "collapsed": True, + "fields": { + "camera": { + "required": False, + "selector": {"boolean": None}, + "description": "Duplicate of a field in another section.", + }, + }, + }, + "extras": { + "collapsed": True, + "fields": {}, + }, + }, + }, + "send_to": { + "description": "Send a message to a target.", + "target": {"entity": {"domain": "notify"}}, + "fields": { + "entity_id": { + "required": False, + "selector": {"number": None}, + "description": "Duplicate of the implicit target argument.", + }, + }, + }, + } + } + + monkeypatch.setattr( + "custom_components.pyscript.stubs.generator.async_get_all_descriptions", fake_service_descriptions + ) + + await pyscript.start() + + stubs_dir = Path(hass.config.path(FOLDER)) / "modules" / "stubs" + generated_target = stubs_dir / "pyscript_generated.py" + stubs_dir.mkdir(parents=True, exist_ok=True) + + try: + await hass.services.async_call( + DOMAIN, SERVICE_GENERATE_STUBS, {}, blocking=True, return_response=True + ) + + generated_content = generated_target.read_text(encoding="utf-8") + # Duplicate field names must not produce duplicate arguments. ast.parse() accepts + # those, only compile() reports them, so compile the generated module. + compile(generated_content, str(generated_target), "exec") + + # The sections themselves must not show up as bogus parameters... + assert "delivery_options" not in generated_content + assert "advanced" not in generated_content + # ...including an empty section... + assert "extras" not in generated_content + # ...a field with a non-section "fields" key is kept... + assert "legacy=None" in generated_content + # ...section fields are flattened, the first definition of a duplicate wins... + assert "subject: str | None=None" in generated_content + assert "camera: str | None=None" in generated_content + assert "camera: bool" not in generated_content + # ...and a field named like the implicit target argument is skipped. + assert "def send_to(*, entity_id: str):" in generated_content + # ...and "multiple" selectors accept a single value or a list. + assert "recipients: str | list[str]" in generated_content + assert "sender: str | list[str]='admin'" in generated_content + assert ( + "channels: Literal['', 'email', 'sms'] | list[Literal['', 'email', 'sms']] | None=None" + in generated_content + ) + finally: + for child in stubs_dir.iterdir(): + child.unlink() + stubs_dir.rmdir()