Add mypy check (#195)

* Add mypy check

* Update mypy path in pipeline.yml

* mypy relate fixes

Author: jan-janssen

.github/workflows/pipeline.yml

@@ -19,6 +19,22 @@ jobs:
         options: "--check --diff"
         src: ./example_workflows/quantum_espresso/qe_xml_parser/src/qe_xml_parser
 
+  mypy:
+    needs: [black]
+    runs-on: ubuntu-latest
+    steps:
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+          architecture: x64
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Install mypy
+        run: pip install mypy
+      - name: Test
+        run: mypy --ignore-missing-imports src/python_workflow_definition
+
   pip_check:
     runs-on: ubuntu-latest
     steps:
@@ -227,4 +243,4 @@ jobs:
     - name: Upload coverage reports to Codecov
       uses: codecov/codecov-action@v5
       with:
-        token: ${{ secrets.CODECOV_TOKEN }}
+        token: ${{ secrets.CODECOV_TOKEN }}

src/python_workflow_definition/aiida.py

@@ -1,6 +1,7 @@
 import traceback
 from dataclasses import replace
 from importlib import import_module
+from typing import Any
 
 from aiida import orm
 from aiida_pythonjob.data.serializer import general_serializer
@@ -90,8 +91,8 @@ def load_workflow_json(file_name: str) -> WorkGraph:
     return wg
 
 
-def write_workflow_json(wg: WorkGraph, file_name: str) -> dict:
-    data = {NODES_LABEL: [], EDGES_LABEL: []}
+def write_workflow_json(wg: WorkGraph, file_name: str) -> None:
+    data: dict[str, Any] = {NODES_LABEL: [], EDGES_LABEL: []}
     node_name_mapping = {}
     data_node_name_mapping = {}
     i = 0

src/python_workflow_definition/cwl/__init__.py

@@ -1,6 +1,7 @@
 import json
 import pickle
 from pathlib import Path
+from typing import Any
 
 from yaml import CDumper as Dumper
 from yaml import dump
@@ -86,7 +87,7 @@ def _write_function_cwl(workflow, directory_path: str = "."):
     export_path.mkdir(parents=True, exist_ok=True)
 
     for i in function_nodes_dict.keys():
-        template = {
+        template: dict[str, Any] = {
             "cwlVersion": "v1.2",
             "class": "CommandLineTool",
             "baseCommand": "python",
@@ -154,7 +155,7 @@ def _write_workflow_config(workflow, directory_path: str = "."):
 
 
 def _write_workflow(workflow, directory_path: str = "."):
-    workflow_template = {
+    workflow_template: dict[str, Any] = {
         "cwlVersion": "v1.2",
         "class": "Workflow",
         "inputs": {},

src/python_workflow_definition/executorlib.py

@@ -1,6 +1,7 @@
 from concurrent.futures import Executor
 from importlib import import_module
 from inspect import isfunction
+from typing import Any
 
 from python_workflow_definition.models import PythonWorkflowDefinitionWorkflow
 from python_workflow_definition.purepython import group_edges, resort_total_lst
@@ -58,7 +59,7 @@ def load_workflow_json(file_name: str, exe: Executor):
     total_lst = group_edges(edges_new_lst)
     total_new_lst = resort_total_lst(total_lst=total_lst, nodes_dict=nodes_new_dict)
 
-    result_dict = {}
+    result_dict: dict[Any, Any] = {}
     last_key = None
     for lst in total_new_lst:
         node = nodes_new_dict[lst[0]]

src/python_workflow_definition/jobflow.py

@@ -1,5 +1,6 @@
 from importlib import import_module
 from inspect import isfunction
+from typing import Any
 
 import numpy as np
 from jobflow import Flow, job
@@ -39,7 +40,7 @@ def _get_nodes_dict(function_dict: dict):
 
 
 def _get_edge_from_dict(
-    target: str, key: str, value_dict: dict, nodes_mapping_dict: dict
+    target: int, key: str, value_dict: dict, nodes_mapping_dict: dict
 ) -> dict:
     if len(value_dict["attributes"]) == 1:
         return {
@@ -199,8 +200,8 @@ def _resort_total_lst(total_dict: dict, nodes_dict: dict) -> dict:
     nodes_without_dep_lst = [
         k for k in nodes_dict.keys() if k not in nodes_with_dep_lst
     ]
-    ordered_lst = []
-    total_new_dict = {}
+    ordered_lst: list = []
+    total_new_dict: dict[Any, dict] = {}
     while len(total_new_dict) < len(total_dict):
         for ind in sorted(total_dict.keys()):
             connect = total_dict[ind]
@@ -215,7 +216,7 @@ def _resort_total_lst(total_dict: dict, nodes_dict: dict) -> dict:
 
 
 def _group_edges(edges_lst: list) -> dict:
-    total_dict = {}
+    total_dict: dict[Any, dict] = {}
     for ed_major in edges_lst:
         target_id = ed_major[TARGET_LABEL]
         tmp_lst = []
@@ -240,7 +241,7 @@ def get_attr_helper(obj, source_handle):
         else:
             return getattr(getattr(obj, "output"), source_handle)
 
-    memory_dict = {}
+    memory_dict: dict[Any, Any] = {}
     for k in total_dict.keys():
         v = nodes_dict[k]
         if isfunction(v):

src/python_workflow_definition/plot.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import networkx as nx
 from IPython.display import SVG, display
 
@@ -25,7 +27,7 @@ def plot(file_name: str):
 
     for edge_tuple in total_lst:
         target_node, edge_dict = edge_tuple
-        edge_label_dict = {}
+        edge_label_dict: dict[Any, list] = {}
         for k, v in edge_dict.items():
             if v[SOURCE_LABEL] not in edge_label_dict:
                 edge_label_dict[v[SOURCE_LABEL]] = []

src/python_workflow_definition/purepython.py

@@ -1,5 +1,6 @@
 from importlib import import_module
 from inspect import isfunction
+from typing import Any
 
 from python_workflow_definition.models import PythonWorkflowDefinitionWorkflow
 from python_workflow_definition.shared import (
@@ -23,7 +24,8 @@ def resort_total_lst(total_lst: list, nodes_dict: dict) -> list:
     nodes_without_dep_lst = [
         k for k in nodes_dict.keys() if k not in nodes_with_dep_lst
     ]
-    ordered_lst, total_new_lst = [], []
+    ordered_lst: list = []
+    total_new_lst: list[list] = []
     while len(total_new_lst) < len(total_lst):
         for ind, connect in total_lst:
             if ind not in ordered_lst:
@@ -86,7 +88,7 @@ def load_workflow_json(file_name: str):
     total_lst = group_edges(edges_new_lst)
     total_new_lst = resort_total_lst(total_lst=total_lst, nodes_dict=nodes_new_dict)
 
-    result_dict = {}
+    result_dict: dict[Any, Any] = {}
     last_key = None
     for lst in total_new_lst:
         node = nodes_new_dict[lst[0]]

src/python_workflow_definition/pyiron_base.py

@@ -1,6 +1,6 @@
 from importlib import import_module
 from inspect import isfunction
-from typing import Optional
+from typing import Any, Optional
 
 import numpy as np
 from pyiron_base import Project, job
@@ -30,7 +30,8 @@ def _resort_total_lst(total_lst: list, nodes_dict: dict) -> list:
     nodes_without_dep_lst = [
         k for k in nodes_dict.keys() if k not in nodes_with_dep_lst
     ]
-    ordered_lst, total_new_lst = [], []
+    ordered_lst: list = []
+    total_new_lst: list[list] = []
     while len(total_new_lst) < len(total_lst):
         for ind, connect in total_lst:
             if ind not in ordered_lst:
@@ -74,7 +75,7 @@ def _get_source(
 def _get_delayed_object_dict(
     total_lst: list, nodes_dict: dict, source_handle_dict: dict, pyiron_project: Project
 ) -> dict:
-    delayed_object_dict = {}
+    delayed_object_dict: dict[Any, DelayedObject] = {}
     for item in total_lst:
         key, input_dict = item
         kwargs = {
@@ -134,8 +135,9 @@ def _get_unique_objects(nodes_dict: dict):
             )
             delayed_object_dict[k]._python_function = get_dict
             delayed_object_dict[k]._input = v
-    unique_lst = []
-    delayed_object_updated_dict, match_dict = {}, {}
+    unique_lst: list = []
+    delayed_object_updated_dict: dict[Any, DelayedObject] = {}
+    match_dict: dict[Any, Any] = {}
     for dobj in delayed_object_dict.keys():
         match = False
         for obj in unique_lst:

src/python_workflow_definition/pyiron_workflow.py

@@ -189,15 +189,15 @@ def write_workflow_json(graph_as_dict: dict, file_name: str = "workflow.json"):
         elif edge[TARGET_LABEL] not in remap_get_list_dict.values():
             edge_get_list_updated_lst.append(edge)
 
-    target_dict = {}
+    target_dict: dict[Any, list] = {}
     for edge in edge_get_list_updated_lst:
         for k in pyiron_workflow_modules.keys():
             if k == edge[TARGET_LABEL]:
                 if k not in target_dict:
                     target_dict[k] = []
                 target_dict[k].append(edge)
 
-    source_dict = {}
+    source_dict: dict[Any, list] = {}
     for edge in edge_get_list_updated_lst:
         for k in pyiron_workflow_modules.keys():
             if k == edge[SOURCE_LABEL]:

src/python_workflow_definition/shared.py

@@ -1,4 +1,5 @@
 from collections import Counter
+from typing import Any
 
 NODES_LABEL = "nodes"
 EDGES_LABEL = "edges"
@@ -31,7 +32,7 @@ def get_kwargs(lst: list) -> dict:
 
 
 def get_source_handles(edges_lst: list) -> dict:
-    source_handle_dict = {}
+    source_handle_dict: dict[Any, list] = {}
     for ed in edges_lst:
         if ed[SOURCE_LABEL] not in source_handle_dict.keys():
             source_handle_dict[ed[SOURCE_LABEL]] = []