add code widget entry point (#487)

This PR adds an entry point for code widgets in the submission step, allowing users to add code widgets for their plugins. 
- Now, only the `pw` code is hardcoded. Other codes, e.g. `dos` and `projwfc`, are from the plugins.
- The related code will be hidden if the workchain property is not selected.

src/aiidalab_qe/app/submission/__init__.py

@@ -16,6 +16,7 @@
 from IPython.display import display
 
 from aiidalab_qe.app.parameters import DEFAULT_PARAMETERS
+from aiidalab_qe.app.utils import get_entry_items
 from aiidalab_qe.common.setup_codes import QESetupWidget
 from aiidalab_qe.workflows import QeAppWorkChain
 
@@ -68,26 +69,30 @@ def __init__(self, qe_auto_setup=True, **kwargs):
         self.pw_code = ComputationalResourcesWidget(
             description="pw.x:", default_calc_job_plugin="quantumespresso.pw"
         )
-        self.dos_code = ComputationalResourcesWidget(
-            description="dos.x:",
-            default_calc_job_plugin="quantumespresso.dos",
-        )
-        self.projwfc_code = ComputationalResourcesWidget(
-            description="projwfc.x:",
-            default_calc_job_plugin="quantumespresso.projwfc",
-        )
 
         self.resources_config = ResourceSelectionWidget()
         self.parallelization = ParallelizationSettings()
 
-        self.set_selected_codes(DEFAULT_PARAMETERS["codes"])
         self.set_resource_defaults()
 
         self.pw_code.observe(self._update_state, "value")
         self.pw_code.observe(self._update_resources, "value")
-        self.dos_code.observe(self._update_state, "value")
-        self.projwfc_code.observe(self._update_state, "value")
 
+        # add plugin's entry points
+        self.codes = {"pw": self.pw_code}
+        self.code_children = [
+            self.codes_title,
+            self.codes_help,
+            self.pw_code,
+        ]
+        self.code_entries = get_entry_items("aiidalab_qe.properties", "code")
+        for _, entry_point in self.code_entries.items():
+            for name, code in entry_point.items():
+                self.codes[name] = code
+                self.code_children.append(self.codes[name])
+        # set default codes
+        self.set_selected_codes(DEFAULT_PARAMETERS["codes"])
+        #
         self.submit_button = ipw.Button(
             description="Submit",
             tooltip="Submit the calculation with the selected parameters.",
@@ -121,11 +126,7 @@ def __init__(self, qe_auto_setup=True, **kwargs):
 
         super().__init__(
             children=[
-                self.codes_title,
-                self.codes_help,
-                self.pw_code,
-                self.dos_code,
-                self.projwfc_code,
+                *self.code_children,
                 self.resources_config,
                 self.parallelization,
                 self.message_area,
@@ -152,47 +153,19 @@ def _identify_submission_blockers(self):
         if self.qe_setup_status.busy or self.sssp_installation_status.busy:
             yield "Background setup processes must finish."
 
-        # No code selected (this is ignored while the setup process is running).
+        # No pw code selected (this is ignored while the setup process is running).
         if self.pw_code.value is None and not self.qe_setup_status.busy:
             yield ("No pw code selected")
-
-        # No code selected for pdos (this is ignored while the setup process is running).
-        if (
-            "pdos" in self.input_parameters.get("workchain", {}).get("properties", [])
-            and (self.dos_code.value is None or self.projwfc_code.value is None)
-            and not self.qe_setup_status.busy
-        ):
-            yield "Calculating the PDOS requires both dos.x and projwfc.x to be set."
-
+        # code related to the selected property is not installed
+        properties = self.input_parameters.get("workchain", {}).get("properties", [])
+        for identifer in properties:
+            for name, code in self.code_entries.get(identifer, {}).items():
+                if code.value is None:
+                    yield f"Calculating the {identifer} property requires code {name} to be set."
         # SSSP library not installed
         if not self.sssp_installation_status.installed:
             yield "The SSSP library is not installed."
 
-        if (
-            "pdos" in self.input_parameters.get("workchain", {}).get("properties", [])
-            and not any(
-                [
-                    self.pw_code.value is None,
-                    self.dos_code.value is None,
-                    self.projwfc_code.value is None,
-                ]
-            )
-            and len(
-                set(
-                    (
-                        orm.load_code(self.pw_code.value).computer.pk,
-                        orm.load_code(self.dos_code.value).computer.pk,
-                        orm.load_code(self.projwfc_code.value).computer.pk,
-                    )
-                )
-            )
-            != 1
-        ):
-            yield (
-                "All selected codes must be installed on the same computer. This is because the "
-                "PDOS calculations rely on large files that are not retrieved by AiiDA."
-            )
-
     def _update_state(self, _=None):
         # If the previous step has failed, this should fail as well.
         if self.previous_step_state is self.State.FAIL:
@@ -225,16 +198,11 @@ def _toggle_install_widgets(self, change):
 
     def _auto_select_code(self, change):
         if change["new"] and not change["old"]:
-            for code in [
-                "pw",
-                "dos",
-                "projwfc",
-            ]:
+            for name, code_widget in self.codes.items():
                 try:
-                    code_widget = getattr(self, f"{code}_code")
                     code_widget.refresh()
                     code_widget.value = orm.load_code(
-                        DEFAULT_PARAMETERS["codes"][code]
+                        DEFAULT_PARAMETERS["codes"][name]
                     ).uuid
                 except NotExistent:
                     pass
@@ -342,7 +310,7 @@ def _observe_state(self, change):
     @tl.observe("previous_step_state")
     def _observe_input_structure(self, _):
         self._update_state()
-        self.set_pdos_status()
+        self.udpate_codes_visibility()
 
     @tl.observe("process")
     def _observe_process(self, change):
@@ -361,11 +329,7 @@ def get_selected_codes(self):
 
         return: A dict with the code names as keys and the code UUIDs as values.
         """
-        codes = {
-            "pw": self.pw_code.value,
-            "dos": self.dos_code.value,
-            "projwfc": self.projwfc_code.value,
-        }
+        codes = {key: code.value for key, code in self.codes.items()}
         return codes
 
     def set_selected_codes(self, codes):
@@ -380,18 +344,21 @@ def _get_code_uuid(code):
                     return None
 
         with self.hold_trait_notifications():
-            # Codes
-            self.pw_code.value = _get_code_uuid(codes["pw"])
-            self.dos_code.value = _get_code_uuid(codes["dos"])
-            self.projwfc_code.value = _get_code_uuid(codes["projwfc"])
-
-    def set_pdos_status(self):
-        if "pdos" in self.input_parameters.get("workchain", {}).get("properties", []):
-            self.dos_code.code_select_dropdown.disabled = False
-            self.projwfc_code.code_select_dropdown.disabled = False
-        else:
-            self.dos_code.code_select_dropdown.disabled = True
-            self.projwfc_code.code_select_dropdown.disabled = True
+            for name, code in self.codes.items():
+                code.value = _get_code_uuid(codes.get(name))
+
+    def udpate_codes_visibility(self):
+        """Hide code if no related property is selected."""
+        # hide all codes except pw
+        for name, code in self.codes.items():
+            if name == "pw":
+                continue
+            code.layout.visibility = "hidden"
+        properties = self.input_parameters.get("workchain", {}).get("properties", [])
+        # show the code if the related property is selected.
+        for identifer in properties:
+            for code in self.code_entries.get(identifer, {}).values():
+                code.layout.visibility = "visible"
 
     def submit(self, _=None):
         """Submit the work chain with the current inputs."""
@@ -477,6 +444,7 @@ def set_submission_parameters(self, parameters):
         self.set_selected_codes(parameters["codes"])
 
     def get_submission_parameters(self):
+        """Get the parameters for the submission step."""
         return {
             "codes": self.get_selected_codes(),
             "resources": self.get_resources(),

src/aiidalab_qe/plugins/bands/workchain.py

@@ -174,7 +174,7 @@ def get_builder(codes, structure, parameters, **kwargs):
     """Get a builder for the PwBandsWorkChain."""
     from copy import deepcopy
 
-    pw_code = codes.get("pw", {})
+    pw_code = codes.get("pw")
     protocol = parameters["workchain"]["protocol"]
     scf_overrides = deepcopy(parameters["advanced"])
     bands_overrides = deepcopy(parameters["advanced"])

src/aiidalab_qe/plugins/pdos/__init__.py

@@ -1,3 +1,5 @@
+from aiidalab_widgets_base import ComputationalResourcesWidget
+
 from aiidalab_qe.common.panel import OutlinePanel
 
 from .result import Result
@@ -10,8 +12,20 @@ class PdosOutline(OutlinePanel):
     help = """"""
 
 
+dos_code = ComputationalResourcesWidget(
+    description="dos.x",
+    default_calc_job_plugin="quantumespresso.dos",
+)
+
+projwfc_code = ComputationalResourcesWidget(
+    description="projwfc.x",
+    default_calc_job_plugin="quantumespresso.projwfc",
+)
+
+
 pdos = {
     "outline": PdosOutline,
+    "code": {"dos": dos_code, "projwfc": projwfc_code},
     "setting": Setting,
     "result": Result,
     "workchain": workchain_and_builder,

src/aiidalab_qe/plugins/pdos/workchain.py

@@ -4,12 +4,40 @@
 PdosWorkChain = WorkflowFactory("quantumespresso.pdos")
 
 
+def check_codes(pw_code, dos_code, projwfc_code):
+    """Check that the codes are installed on the same computer."""
+    if (
+        not any(
+            [
+                pw_code is None,
+                dos_code is None,
+                projwfc_code is None,
+            ]
+        )
+        and len(
+            set(
+                (
+                    pw_code.computer.pk,
+                    dos_code.computer.pk,
+                    projwfc_code.computer.pk,
+                )
+            )
+        )
+        != 1
+    ):
+        raise ValueError(
+            "All selected codes must be installed on the same computer. This is because the "
+            "PDOS calculations rely on large files that are not retrieved by AiiDA."
+        )
+
+
 def get_builder(codes, structure, parameters, **kwargs):
     from copy import deepcopy
 
-    pw_code = codes.get("pw", None)
-    dos_code = codes.get("dos", None)
-    projwfc_code = codes.get("projwfc", None)
+    pw_code = codes.get("pw")
+    dos_code = codes.get("dos")
+    projwfc_code = codes.get("projwfc")
+    check_codes(pw_code, dos_code, projwfc_code)
     protocol = parameters["workchain"]["protocol"]
 
     scf_overrides = deepcopy(parameters["advanced"])

tests/conftest.py

@@ -299,19 +299,17 @@ def app(pw_code, dos_code, projwfc_code):
     app = App(qe_auto_setup=False)
     # set up codes
     app.submit_step.pw_code.refresh()
-    app.submit_step.dos_code.refresh()
-    app.submit_step.projwfc_code.refresh()
+    app.submit_step.codes["dos"].refresh()
+    app.submit_step.codes["projwfc"].refresh()
     app.submit_step.pw_code.value = (
         app.submit_step.pw_code.code_select_dropdown.options[pw_code.full_label]
     )
-    app.submit_step.dos_code.value = (
-        app.submit_step.dos_code.code_select_dropdown.options[dos_code.full_label]
-    )
-    app.submit_step.projwfc_code.value = (
-        app.submit_step.projwfc_code.code_select_dropdown.options[
-            projwfc_code.full_label
-        ]
-    )
+    app.submit_step.codes["dos"].value = app.submit_step.codes[
+        "dos"
+    ].code_select_dropdown.options[dos_code.full_label]
+    app.submit_step.codes["projwfc"].value = app.submit_step.codes[
+        "projwfc"
+    ].code_select_dropdown.options[projwfc_code.full_label]
 
     yield app
 

tests/test_codes.py

@@ -0,0 +1,49 @@
+def test_code_not_selected(submit_app_generator):
+    """Test if there is an error when the code is not selected."""
+    app = submit_app_generator()
+    app.submit_step.codes["dos"].value = None
+    app.submit_step._create_builder()
+
+
+def test_set_selected_codes(submit_app_generator):
+    """Test set_selected_codes method."""
+    from aiidalab_qe.app.submission import SubmitQeAppWorkChainStep
+
+    app = submit_app_generator()
+    submit_step = app.submit_step
+
+    submit_step._create_builder()
+
+    new_submit_step = SubmitQeAppWorkChainStep(qe_auto_setup=False)
+    new_submit_step.set_selected_codes(submit_step.ui_parameters["codes"])
+
+    assert new_submit_step.get_selected_codes() == submit_step.get_selected_codes()
+
+
+def test_udpate_codes_visibility():
+    """Test udpate_codes_visibility method.
+    If the workchain property is not selected, the related code should be hidden.
+    """
+    from aiidalab_qe.app.submission import SubmitQeAppWorkChainStep
+
+    submit = SubmitQeAppWorkChainStep(qe_auto_setup=False)
+    submit.udpate_codes_visibility()
+    assert submit.codes["dos"].layout.visibility == "hidden"
+    submit.input_parameters = {"workchain": {"properties": ["pdos"]}}
+    submit.udpate_codes_visibility()
+    assert submit.codes["dos"].layout.visibility == "visible"
+
+
+def test_identify_submission_blockers(app):
+    """Test identify_submission_blockers method."""
+    submit = app.submit_step
+    blockers = list(submit._identify_submission_blockers())
+    # there is one blocker: ['The SSSP library is not installed.']
+    assert len(blockers) == 1
+    submit.input_parameters = {"workchain": {"properties": ["pdos"]}}
+    blockers = list(submit._identify_submission_blockers())
+    assert len(blockers) == 1
+    # set dos code to None, will introduce another blocker
+    submit.codes["dos"].value = None
+    blockers = list(submit._identify_submission_blockers())
+    assert len(blockers) == 2

tests/test_submit_qe_workchain.py

@@ -1,27 +1,3 @@
-def test_code_not_selected(submit_app_generator):
-    """Test if there is an error when the code is not selected."""
-    app = submit_app_generator()
-    app.submit_step.dos_code.value = None
-    app.submit_step._create_builder()
-
-
-def test_reload_selected_code(submit_app_generator):
-    """Test set_selected_codes method."""
-    from aiidalab_qe.app.submission import SubmitQeAppWorkChainStep
-
-    app = submit_app_generator()
-    submit_step = app.submit_step
-
-    submit_step._create_builder()
-
-    new_submit_step = SubmitQeAppWorkChainStep(qe_auto_setup=False)
-    new_submit_step.set_selected_codes(submit_step.ui_parameters["codes"])
-
-    assert new_submit_step.pw_code.value == submit_step.pw_code.value
-    assert new_submit_step.dos_code.value == submit_step.dos_code.value
-    assert new_submit_step.projwfc_code.value == submit_step.projwfc_code.value
-
-
 def test_create_builder_default(
     data_regression,
     submit_app_generator,