WIP: to mypy

Author: unkcpz

.pre-commit-config.yaml

@@ -66,4 +66,14 @@ repos:
                   docker/.*|
                   docs/.*|
 
+                  tests/.*|
+                  tests_integration/.*|
+
+                  # Exclude the following files temporarily.
+                  src/aiidalab_qe/app/submission/__init__.py|
+                  src/aiidalab_qe/app/parameters/__init__.py|
+                  src/aiidalab_qe/__main__.py|
+                  src/aiidalab_qe/common/.*|
+                  src/aiidalab_qe/plugins/.*|
+
                 )$

setup.cfg

@@ -46,6 +46,8 @@ dev =
     pytest-regressions~=2.2
     pgtest==1.3.1
     pytest-cov~=4.0
+    beautifulsoup4~=4.10
+    types-beautifulsoup4~=4.10
 
 [options.package_data]
 aiidalab_qe.app.parameters = qeapp.yaml
@@ -71,7 +73,25 @@ show_traceback = true
 strict = true
 warn_return_any = false
 warn_unreachable = true
-explicit_package_bases = true
+no_namespace_packages = true
+
+[mypy-ipywidgets.*]
+ignore_missing_imports = true
+
+[mypy-aiidalab_widgets_base.*]
+ignore_missing_imports = true
+
+[mypy-aiida_quantumespresso.*]
+ignore_missing_imports = true
+
+[mypy-IPython.*]
+ignore_missing_imports = true
+
+[mypy-setuptools.*]
+ignore_missing_imports = true
+
+[mypy-aiida_pseudo.*]
+ignore_missing_imports = true
 
 [aiidalab]
 title = Quantum ESPRESSO

src/aiidalab_qe/app/configuration/advanced.py

@@ -7,6 +7,7 @@
 
 import ipywidgets as ipw
 import traitlets as tl
+import typing as t
 from aiida import orm
 from aiida_quantumespresso.calculations.functions.create_kpoints_from_distance import (
     create_kpoints_from_distance,
@@ -222,7 +223,7 @@ def update_settings(self, **kwargs):
     def get_panel_value(self):
         # create the the initial_magnetic_moments as None (Default)
         # XXX: start from parameters = {} and then bundle the settings by purposes (e.g. pw, bands, etc.)
-        parameters = {
+        parameters: dict[str, t.Any] = {
             "initial_magnetic_moments": None,
             "pw": {
                 "parameters": {

src/aiidalab_qe/app/configuration/pseudos.py

@@ -3,6 +3,7 @@
 
 import io
 import re
+import typing as t
 
 import ipywidgets as ipw
 import traitlets as tl
@@ -19,6 +20,9 @@
     PseudoFamily,
 )
 
+if t.TYPE_CHECKING:
+    from aiida_pseudo.data.pseudo.upf import UpfData as TypeUpfData
+
 UpfData = DataFactory("pseudo.upf")
 SsspFamily = GroupFactory("pseudo.family.sssp")
 PseudoDojoFamily = GroupFactory("pseudo.family.pseudo_dojo")
@@ -174,9 +178,8 @@ def set_text_color(self, change):
             self.dft_functional_help,
             self.dft_functional_prompt,
         ):
-            old_opacity = re.match(
-                r"[\s\S]+opacity:([\S]+);[\S\s]+", html.value
-            ).groups()[0]
+            matched = re.match(r"[\s\S]+opacity:([\S]+);[\S\s]+", html.value)
+            old_opacity = matched.group(1) if matched else 1.0
             html.value = html.value.replace(
                 f"opacity:{old_opacity};", f"opacity:{opacity};"
             )
@@ -344,7 +347,7 @@ def _reset(self):
             return
 
         try:
-            pseudos = pseudo_family.get_pseudos(structure=self.structure)
+            pseudos = pseudo_family.get_pseudos(structure=self.structure)  # type: ignore
             # get cutoffs dict of all elements
             cutoffs = self._get_cutoffs(pseudo_family)
         except ValueError as exception:
@@ -356,23 +359,23 @@ def _reset(self):
         self.pseudos = {kind: pseudo.uuid for kind, pseudo in pseudos.items()}
         self.set_pseudos(self.pseudos, cutoffs)
 
-    def _get_pseudos_family(self, pseudo_family: str) -> orm.Group:
+    def _get_pseudos_family(self, pseudo_family: str) -> t.Optional[orm.Group]:
         """Get the pseudo family from the database."""
         try:
             pseudo_set = (PseudoDojoFamily, SsspFamily, CutoffsPseudoPotentialFamily)
-            pseudo_family = (
+            pseudo_family_node = (
                 orm.QueryBuilder()
                 .append(pseudo_set, filters={"label": pseudo_family})
                 .one()[0]
             )
+            return pseudo_family_node
+
         except exceptions.NotExistent as exception:
             raise exceptions.NotExistent(
                 f"required pseudo family `{pseudo_family}` is not installed. Please use `aiida-pseudo install` to"
                 "install it."
             ) from exception
 
-        return pseudo_family
-
     def _get_cutoffs(self, pseudo_family):
         """Get the cutoffs from the pseudo family."""
         from aiida_pseudo.common.units import U
@@ -462,7 +465,7 @@ class PseudoUploadWidget(ipw.HBox):
     def __init__(
         self,
         kind: str = "",
-        pseudo: UpfData | None = None,
+        pseudo: TypeUpfData | None = None,
         cutoffs: dict | None = None,
         **kwargs,
     ):

src/aiidalab_qe/app/result/__init__.py

@@ -103,7 +103,7 @@ def _on_click_kill_button(self, _=None):
         First kill the process, then update the kill button layout.
         """
         workchain = [orm.load_node(self.process)]
-        control.kill_processes(workchain)
+        control.kill_processes(workchain)  # type: ignore
 
         # update the kill button layout
         self._update_kill_button_layout()

src/aiidalab_qe/app/result/workchain_viewer.py

@@ -167,11 +167,10 @@ def __init__(self, node, export_dir=None, **kwargs):
         self._download_archive_button.on_click(self._download_archive)
         self._download_button_container = ipw.Box([self._download_archive_button])
 
-        if node.exit_status != 0:
+        if node.exit_status != 0 and (final_calcjob := self._get_final_calcjob(node)):
             title = ipw.HTML(
                 f"<h4>Workflow failed with exit status [{ node.exit_status }]</h4>"
             )
-            final_calcjob = self._get_final_calcjob(node)
             env = Environment()
             template = resources.read_text(static, "workflow_failure.jinja")
             style = resources.read_text(static, "style.css")
@@ -224,7 +223,9 @@ def _download_archive(self, _):
                     # Create archive file.
                     with TemporaryDirectory() as tmpdir:
                         self._prepare_calcjob_io(self.node, Path(tmpdir))
-                        shutil.make_archive(fn_archive.with_suffix(""), "zip", tmpdir)
+                        shutil.make_archive(
+                            str(fn_archive.with_suffix("")), "zip", tmpdir
+                        )
             Path(fn_lockfile).unlink()  # Delete lock file.
         except Timeout:
             # Failed to obtain lock, presuming some other process is working on it.

src/aiidalab_qe/app/submission/__init__.py

@@ -191,7 +191,7 @@ def _update_state(self, _=None):
 
     def _toggle_install_widgets(self, change):
         if change["new"]:
-            self.children = [
+            self.children = [  # type: ignore
                 child for child in self.children if child is not change["owner"]
             ]
 
@@ -224,8 +224,8 @@ def _show_alert_message(self, message, alert_class="info"):
     def _update_resources(self, change):
         if change["new"] and (
             change["old"] is None
-            or orm.load_code(change["new"]).computer.pk
-            != orm.load_code(change["old"]).computer.pk
+            or orm.load_code(change["new"]).computer.pk  # type: ignore
+            != orm.load_code(change["old"]).computer.pk  # type: ignore
         ):
             self.set_resource_defaults(orm.load_code(change["new"]).computer)
 
@@ -277,7 +277,7 @@ def _check_resources(self):
 
         num_cpus = self.resources_config.num_cpus.value
         on_localhost = (
-            orm.load_node(self.pw_code.value).computer.hostname == "localhost"
+            orm.load_node(self.pw_code.value).computer.hostname == "localhost"  # type: ignore
         )
         if self.pw_code.value and on_localhost and num_cpus > 1:
             self._show_alert_message(
@@ -376,7 +376,7 @@ def submit(self, _=None):
 
         self._update_state()
 
-    def _generate_label(self) -> dict:
+    def _generate_label(self) -> str:
         """Generate a label for the work chain based on the input parameters."""
         formula = self.input_structure.get_formula()
         properties = [
@@ -392,8 +392,7 @@ def _generate_label(self) -> dict:
         else:
             properties_info = f"properties on {', '.join(properties)}"
 
-        label = "{} {} {}".format(formula, relax_info, properties_info)
-        return label
+        return f"{formula} {relax_info} {properties_info}"
 
     def _create_builder(self) -> ProcessBuilderNamespace:
         """Create the builder for the `QeAppWorkChain` submit."""

src/aiidalab_qe/common/panel.py

@@ -5,9 +5,11 @@
 
     AiiDAlab Team
 """
+from __future__ import annotations
+
 import ipywidgets as ipw
 
-DEFAULT_PARAMETERS = {}
+DEFAULT_PARAMETERS: dict = {}
 
 
 class Panel(ipw.VBox):
@@ -100,7 +102,7 @@ class ResultPanel(Panel):
 
     title = "Result"
     # to specify which plugins (outputs) are needed for this result panel.
-    workchain_labels = []
+    workchain_labels: list = []
 
     def __init__(self, node=None, **kwargs):
         self.node = node

src/aiidalab_qe/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561

src/aiidalab_qe/workflows/__init__.py

@@ -84,9 +84,9 @@ def define(cls, spec):
             i += 1
         spec.outline(
             cls.setup,
-            if_(cls.should_run_relax)(
-                cls.run_relax,
-                cls.inspect_relax
+            if_(cls.should_run_relax)( # type: ignore
+                cls.run_relax,  # type: ignore
+                cls.inspect_relax   # type: ignore
             ),
             cls.run_plugin,
             cls.inspect_plugin,