Merge pull request #21 from fleXible-public/feat/pythonic_code

refactor: pythonic code style, tests for Function class

Author: craigbarratt

custom_components/pyscript/function.py

@@ -143,49 +143,39 @@ async def service_call(cls, domain, name, **kwargs):
     async def service_completions(cls, root):
         """Return possible completions of HASS services."""
         words = set()
-        services = await async_get_all_descriptions(cls.hass)
+        services = cls.hass.services.async_services()
         num_period = root.count(".")
         if num_period == 1:
-            domain, srv_root = root.split(".")
+            domain, svc_root = root.split(".")
             if domain in services:
-                for srv in services[domain].keys():
-                    if srv.lower().startswith(srv_root):
-                        words.add(f"{domain}.{srv}")
+                words |= {f"{domain}.{svc}" for svc in services[domain] if svc.lower().startswith(svc_root)}
         elif num_period == 0:
-            for domain in services.keys():
-                if domain.lower().startswith(root):
-                    words.add(domain)
+            words |= {domain for domain in services if domain.lower().startswith(root)}
+
         return words
 
     @classmethod
     async def func_completions(cls, root):
         """Return possible completions of functions."""
-        words = set()
-        funcs = cls.functions.copy()
-        funcs.update(cls.ast_functions)
-        for name in funcs.keys():
-            if name.lower().startswith(root):
-                words.add(name)
+        funcs = {**cls.functions, **cls.ast_functions}
+        words = {name for name in funcs if name.lower().startswith(root)}
+
         return words
 
     @classmethod
     def register(cls, funcs):
         """Register functions to be available for calling."""
-        for name, func in funcs.items():
-            cls.functions[name] = func
+        cls.functions.update(funcs)
 
     @classmethod
     def register_ast(cls, funcs):
         """Register functions that need ast context to be available for calling."""
-        for name, func in funcs.items():
-            cls.ast_functions[name] = func
+        cls.ast_functions.update(funcs)
 
     @classmethod
     def install_ast_funcs(cls, ast_ctx):
         """Install ast functions into the local symbol table."""
-        sym_table = {}
-        for name, func in cls.ast_functions.items():
-            sym_table[name] = func(ast_ctx)
+        sym_table = {name: func(ast_ctx) for name, func in cls.ast_functions.items()}
         ast_ctx.set_local_sym_table(sym_table)
 
     @classmethod
@@ -194,17 +184,20 @@ def get(cls, name):
         func = cls.functions.get(name, None)
         if func:
             return func
-        parts = name.split(".", 1)
-        if len(parts) != 2:
+
+        name_parts = name.split(".")
+        if len(name_parts) != 2:
             return None
-        domain = parts[0]
-        service = parts[1]
-        if not cls.hass.services.has_service(domain, service):
+
+        domain, service = name_parts
+        if not cls.service_has_service(domain, service):
             return None
 
         async def service_call(*args, **kwargs):
             await cls.hass.services.async_call(domain, service, kwargs)
 
+        # service_call = functools.partial(cls.service_call, domain, service)
+
         return service_call
 
     @classmethod
@@ -218,17 +211,14 @@ async def run_coro(cls, coro):
         try:
             await coro
         except asyncio.CancelledError:
-            if task in cls.unique_task2name:
-                cls.unique_name2task.pop(cls.unique_task2name[task], None)
-                cls.unique_task2name.pop(task, None)
-            cls.our_tasks.discard(task)
             raise
         except Exception:  # pylint: disable=broad-except
             _LOGGER.error("run_coro: %s", traceback.format_exc(-1))
-        if task in cls.unique_task2name:
-            cls.unique_name2task.pop(cls.unique_task2name[task], None)
-            cls.unique_task2name.pop(task, None)
-        cls.our_tasks.discard(task)
+        finally:
+            if task in cls.unique_task2name:
+                del cls.unique_name2task[cls.unique_task2name[task]]
+                del cls.unique_task2name[task]
+            cls.our_tasks.discard(task)
 
     @classmethod
     def create_task(cls, coro):

tests/test_func.py tests/test_function.pytests/test_func.py renamed to tests/test_function.py

@@ -1,18 +1,102 @@
 """Test the pyscript component."""
-from ast import literal_eval
 import asyncio
-from datetime import datetime as dt
 import pathlib
 import time
+from ast import literal_eval
+from datetime import datetime as dt
 
-from custom_components.pyscript.const import DOMAIN
-import custom_components.pyscript.trigger as trigger
-
+import pytest
 from homeassistant import loader
 from homeassistant.const import EVENT_HOMEASSISTANT_STARTED, EVENT_STATE_CHANGED
 from homeassistant.setup import async_setup_component
+from pytest_homeassistant.async_mock import MagicMock, Mock, mock_open, patch
+
+import custom_components.pyscript.trigger as trigger
+from custom_components.pyscript.const import DOMAIN
+from custom_components.pyscript.function import Function
+
+
+@pytest.fixture()
+def ast_functions():
+    return {
+        "domain_ast.func_name": lambda ast_ctx: ast_ctx.func(),
+        "domain_ast.other_func": lambda ast_ctx: ast_ctx.func(),
+    }
+
+
+@pytest.fixture
+def functions():
+    mock_func = Mock()
+
+    return {
+        "domain.func_1": mock_func,
+        "domain.func_2": mock_func,
+        "helpers.get_today": mock_func,
+        "helpers.entity_id": mock_func,
+    }
+
 
-from pytest_homeassistant.async_mock import mock_open, patch
+@pytest.fixture
+def services():
+    return {
+        "domain": {"turn_on": None, "turn_off": None, "toggle": None},
+        "helpers": {"set_state": None, "restart": None},
+    }
+
+
+def test_install_ast_funcs(ast_functions):
+    ast_ctx = MagicMock()
+    ast_ctx.func.return_value = "ok"
+
+    with patch.object(Function, "ast_functions", ast_functions):
+        Function.install_ast_funcs(ast_ctx)
+        assert len(ast_ctx.method_calls) == 3
+
+
+@pytest.mark.parametrize(
+    "root,expected",
+    [
+        ("helpers", {"helpers.entity_id", "helpers.get_today"}),
+        (
+                "domain",
+                {
+                    "domain.func_2",
+                    "domain_ast.func_name",
+                    "domain_ast.other_func",
+                    "domain.func_1",
+                },
+        ),
+        ("domain_", {"domain_ast.func_name", "domain_ast.other_func"}),
+        ("domain_ast.func", {"domain_ast.func_name"}),
+        ("no match", set()),
+    ],
+    ids=lambda x: x if not isinstance(x, (set,)) else f"set({len(x)})",
+)
+async def test_func_completions(ast_functions, functions, root, expected):
+    with patch.object(Function, "ast_functions", ast_functions), patch.object(
+            Function, "functions", functions
+    ):
+        words = await Function.func_completions(root)
+        assert words == expected
+
+
+@pytest.mark.parametrize(
+    "root,expected",
+    [
+        ("do", {"domain"}),
+        ("domain.t", {"domain.toggle", "domain.turn_on", "domain.turn_off"}),
+        ("domain.turn", {"domain.turn_on", "domain.turn_off"}),
+        ("helpers.set", {"helpers.set_state"}),
+        ("no match", set()),
+    ],
+    ids=lambda x: x if not isinstance(x, (set,)) else f"set({len(x)})",
+)
+async def test_service_completions(root, expected, hass, services):
+    with patch.object(
+            hass.services, "async_services", return_value=services
+    ), patch.object(Function, "hass", hass):
+        words = await Function.service_completions(root)
+        assert words == expected
 
 
 async def setup_script(hass, notify_q, now, source):
@@ -34,14 +118,10 @@ async def setup_script(hass, notify_q, now, source):
 
     with patch(
         "homeassistant.loader.async_get_integration", return_value=integration,
-    ), patch(
-        "custom_components.pyscript.os.path.isdir", return_value=True
-    ), patch(
+    ), patch("custom_components.pyscript.os.path.isdir", return_value=True), patch(
         "custom_components.pyscript.glob.iglob", return_value=scripts
     ), patch(
-        "custom_components.pyscript.open",
-        mock_open(read_data=source),
-        create=True,
+        "custom_components.pyscript.open", mock_open(read_data=source), create=True,
     ), patch(
         "custom_components.pyscript.trigger.dt_now", return_value=now
     ):
@@ -157,7 +237,8 @@ def func4(trigger_type=None, event_type=None, **kwargs):
     seq_num += 1
     res = task.wait_until(state_trigger="pyscript.f4var2 == '10'", timeout=10)
     log.info(f"func4 trigger_type = {res}")
-    pyscript.done = [seq_num, res, pyscript.setVar1, pyscript.setVar1.attr1, state.get("pyscript.setVar1.attr2"), pyscript.setVar2, state.get("pyscript.setVar3")]
+    pyscript.done = [seq_num, res, pyscript.setVar1, pyscript.setVar1.attr1, state.get("pyscript.setVar1.attr2"), 
+    pyscript.setVar2, state.get("pyscript.setVar3")]
 
     seq_num += 1
     #