support reloading allow all imports through reload service and config entry options UI

Author: raman325

custom_components/pyscript/__init__.py

@@ -65,7 +65,7 @@ async def async_setup_entry(hass, config_entry):
         await hass.async_add_executor_job(os.makedirs, pyscript_folder)
 
     hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][CONF_ALLOW_ALL_IMPORTS] = config_entry.data.get(CONF_ALLOW_ALL_IMPORTS)
+    hass.data[DOMAIN] = config_entry
 
     State.set_pyscript_config(config_entry.data)
 

custom_components/pyscript/config_flow.py

@@ -5,7 +5,8 @@
 import voluptuous as vol
 
 from homeassistant import config_entries
-from homeassistant.config_entries import SOURCE_IMPORT
+from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
+from homeassistant.core import callback
 
 from .const import CONF_ALLOW_ALL_IMPORTS, DOMAIN
 
@@ -14,12 +15,52 @@
 )
 
 
+class PyscriptOptionsConfigFlow(config_entries.OptionsFlow):
+    """Handle a pyscript options flow."""
+
+    def __init__(self, config_entry: ConfigEntry) -> None:
+        """Initialize pyscript options flow."""
+        self.config_entry = config_entry
+
+    async def async_step_init(self, user_input: Dict[str, Any] = None) -> Dict[str, Any]:
+        """Manage the pyscript options."""
+        if self.config_entry.source == SOURCE_IMPORT:
+            return self.async_abort(reason="no_ui_configuration_allowed")
+
+        if user_input is None:
+            return self.async_show_form(
+                step_id="init",
+                data_schema=vol.Schema(
+                    {
+                        vol.Optional(
+                            CONF_ALLOW_ALL_IMPORTS, default=self.config_entry.data[CONF_ALLOW_ALL_IMPORTS],
+                        ): bool
+                    },
+                    extra=vol.ALLOW_EXTRA,
+                ),
+            )
+
+        if user_input[CONF_ALLOW_ALL_IMPORTS] != self.config_entry.data[CONF_ALLOW_ALL_IMPORTS]:
+            updated_data = self.config_entry.data.copy()
+            updated_data[CONF_ALLOW_ALL_IMPORTS] = user_input[CONF_ALLOW_ALL_IMPORTS]
+            self.hass.config_entries.async_update_entry(entry=self.config_entry, data=updated_data)
+            return self.async_create_entry(title="", data=user_input)
+
+        return self.async_abort(reason="no_update")
+
+
 class PyscriptConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     """Handle a pyscript config flow."""
 
     VERSION = 1
     CONNECTION_CLASS = config_entries.CONN_CLASS_LOCAL_PUSH
 
+    @staticmethod
+    @callback
+    def async_get_options_flow(config_entry: ConfigEntry) -> PyscriptOptionsConfigFlow:
+        """Get the options flow for this handler."""
+        return PyscriptOptionsConfigFlow(config_entry)
+
     async def async_step_user(self, user_input: Dict[str, Any] = None) -> Dict[str, Any]:
         """Handle a flow initialized by the user."""
         if user_input is not None:

custom_components/pyscript/eval.py

@@ -16,7 +16,7 @@
 from homeassistant.const import SERVICE_RELOAD
 from homeassistant.helpers.service import async_set_service_schema
 
-from .const import ALLOWED_IMPORTS, DOMAIN, LOGGER_PATH, SERVICE_JUPYTER_KERNEL_START
+from .const import ALLOWED_IMPORTS, CONF_ALLOW_ALL_IMPORTS, DOMAIN, LOGGER_PATH, SERVICE_JUPYTER_KERNEL_START
 from .function import Function
 from .state import State
 
@@ -721,7 +721,7 @@ def __init__(self, name, global_ctx, logger_name=None):
         self.logger_handlers = set()
         self.logger = None
         self.set_logger_name(logger_name if logger_name is not None else self.name)
-        self.allow_all_imports = Function.hass.data.get(DOMAIN, {}).get("allow_all_imports", False)
+        self.config_entry = Function.hass.data.get(DOMAIN, {})
 
     async def ast_not_implemented(self, arg, *args):
         """Raise NotImplementedError exception for unimplemented AST types."""
@@ -769,7 +769,10 @@ async def ast_import(self, arg):
                 self.exception_long = error_ctx.exception_long
                 raise self.exception_obj
             if not mod:
-                if not self.allow_all_imports and imp.name not in ALLOWED_IMPORTS:
+                if (
+                    not self.config_entry.data.get(CONF_ALLOW_ALL_IMPORTS, False)
+                    and imp.name not in ALLOWED_IMPORTS
+                ):
                     raise ModuleNotFoundError(f"import of {imp.name} not allowed")
                 if imp.name not in sys.modules:
                     mod = await Function.hass.async_add_executor_job(importlib.import_module, imp.name)
@@ -799,7 +802,10 @@ async def ast_importfrom(self, arg):
             self.exception_long = error_ctx.exception_long
             raise self.exception_obj
         if not mod:
-            if not self.allow_all_imports and arg.module not in ALLOWED_IMPORTS:
+            if (
+                not self.config_entry.data.get(CONF_ALLOW_ALL_IMPORTS, False)
+                and arg.module not in ALLOWED_IMPORTS
+            ):
                 raise ModuleNotFoundError(f"import from {arg.module} not allowed")
             if arg.module not in sys.modules:
                 mod = await Function.hass.async_add_executor_job(importlib.import_module, arg.module)

custom_components/pyscript/strings.json

@@ -14,5 +14,19 @@
       "single_instance_allowed": "Already configured. Only a single configuration possible.",
       "updated_entry": "This entry has already been setup but the configuration has been updated."
     }
+  },
+  "options": {
+    "step": {
+      "init": {
+        "title": "Update pyscript configuration",
+        "data": {
+          "allow_all_imports": "Allow All Imports?"
+        }
+      }
+    },
+    "abort": {
+      "no_ui_configuration_allowed": "This entry was created via `configuration.yaml`, so all configuration parameters must be updated there. The [`pyscript.reload`](developer-tools/service) service will allow you to apply the changes you make to `configuration.yaml` without restarting your Home Assistant instance.",
+      "no_update": "There was nothing to update."
+    }
   }
-}
+}

custom_components/pyscript/translations/en.json

@@ -14,5 +14,19 @@
       "single_instance_allowed": "Already configured. Only a single configuration possible.",
       "updated_entry": "This entry has already been setup but the configuration has been updated."
     }
+  },
+  "options": {
+    "step": {
+      "init": {
+        "title": "Update pyscript configuration",
+        "data": {
+          "allow_all_imports": "Allow All Imports?"
+        }
+      }
+    },
+    "abort": {
+      "no_ui_configuration_allowed": "This entry was created via `configuration.yaml`, so all configuration parameters must be updated there. The [`pyscript.reload`](developer-tools/service) service will allow you to apply the changes you make to `configuration.yaml` without restarting your Home Assistant instance.",
+      "no_update": "There is nothing to update."
+    }
   }
-}
+}

tests/test_config_flow.py

@@ -174,3 +174,65 @@ async def test_import_flow_update_import(hass):
     assert result["reason"] == "updated_entry"
 
     assert hass.config_entries.async_entries(DOMAIN)[0].data == {"apps": {"test_app": {"param": 1}}}
+
+
+async def test_options_flow_import(hass):
+    """Test options flow aborts because configuration needs to be managed via configuration.yaml."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": SOURCE_IMPORT}, data=PYSCRIPT_SCHEMA({CONF_ALLOW_ALL_IMPORTS: True})
+    )
+    await hass.async_block_till_done()
+    assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
+    entry = result["result"]
+
+    await hass.config_entries.options.async_init(entry.entry_id, data=None)
+
+    assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT
+    assert result["reason"] == "no_ui_configuration_allowed"
+
+
+async def test_options_flow_user_change(hass):
+    """Test options flow updates config entry when options change."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": SOURCE_USER}, data=PYSCRIPT_SCHEMA({CONF_ALLOW_ALL_IMPORTS: True})
+    )
+    await hass.async_block_till_done()
+    assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
+    entry = result["result"]
+
+    result = await hass.config_entries.options.async_init(entry.entry_id)
+
+    assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
+    assert result["step_id"] == "init"
+
+    result = await hass.config_entries.options.async_configure(
+        result["flow_id"], user_input={CONF_ALLOW_ALL_IMPORTS: False}
+    )
+
+    assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
+    assert result["title"] == ""
+    assert result["data"][CONF_ALLOW_ALL_IMPORTS] is False
+
+    assert entry.data[CONF_ALLOW_ALL_IMPORTS] is False
+
+
+async def test_options_flow_user_no_change(hass):
+    """Test options flow aborts when options don't change."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": SOURCE_USER}, data=PYSCRIPT_SCHEMA({CONF_ALLOW_ALL_IMPORTS: True})
+    )
+    await hass.async_block_till_done()
+    assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
+    entry = result["result"]
+
+    result = await hass.config_entries.options.async_init(entry.entry_id)
+
+    assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
+    assert result["step_id"] == "init"
+
+    result = await hass.config_entries.options.async_configure(
+        result["flow_id"], user_input={CONF_ALLOW_ALL_IMPORTS: True}
+    )
+
+    assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT
+    assert result["reason"] == "no_update"

tests/test_unit_eval.py

@@ -1,9 +1,11 @@
 """Unit tests for Python interpreter."""
 
+from custom_components.pyscript.const import CONF_ALLOW_ALL_IMPORTS, DOMAIN
 from custom_components.pyscript.eval import AstEval
 from custom_components.pyscript.function import Function
 from custom_components.pyscript.global_ctx import GlobalContext, GlobalContextMgr
 from custom_components.pyscript.state import State
+from pytest_homeassistant_custom_component.common import MockConfigEntry
 
 evalTests = [
     ["1", 1],
@@ -882,6 +884,7 @@ async def run_one_test(test_data):
 
 async def test_eval(hass):
     """Test interpreter."""
+    hass.data[DOMAIN] = MockConfigEntry(domain=DOMAIN, data={CONF_ALLOW_ALL_IMPORTS: False})
     Function.init(hass)
     State.init(hass)
     State.register_functions()
@@ -1062,6 +1065,7 @@ async def run_one_test_exception(test_data):
 
 async def test_eval_exceptions(hass):
     """Test interpreter exceptions."""
+    hass.data[DOMAIN] = MockConfigEntry(domain=DOMAIN, data={CONF_ALLOW_ALL_IMPORTS: False})
     Function.init(hass)
     State.init(hass)
     State.register_functions()