diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index b28d5c75..d4564ca9 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -83,3 +83,21 @@ jobs:
         verdi presto
     - name: Run formatter and linter
       run: hatch fmt --check
+
+  typechecking:
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.12
+      uses: actions/setup-python@v5
+      with:
+        python-version: 3.12
+    - name: Install hatch
+      run: |
+        pip install --upgrade pip
+        pip install hatch
+    - name: Install Graphviz
+      run: sudo apt-get install graphviz graphviz-dev
+    - name: Run formatter and linter
+      run: hatch run types:check
diff --git a/pyproject.toml b/pyproject.toml
index 7f9a4c6a..873202dc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -117,3 +117,30 @@ serve = [
 deploy = [
   "mkdocs gh-deploy --no-history -f docs/mkdocs.yml"
 ]
+
+[tool.hatch.envs.types]
+python = "3.12"
+extra-dependencies = [
+  "mypy>=1.0.0",
+  "pytest",
+  "lxml-stubs",
+  "types-setuptools",
+  "types-docutils",
+  "types-colorama",
+  "types-Pygments"
+]
+
+[tool.hatch.envs.types.scripts]
+check = "mypy --no-incremental {args:.}"
+
+[[tool.mypy.overrides]]
+module = ["isoduration", "isoduration.*"]
+follow_untyped_imports = true
+
+[[tool.mypy.overrides]]
+module = ["pygraphviz"]
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = ["f90nml"]
+ignore_missing_imports = true
diff --git a/src/sirocco/core/__init__.py b/src/sirocco/core/__init__.py
index e2c15e41..1da21226 100644
--- a/src/sirocco/core/__init__.py
+++ b/src/sirocco/core/__init__.py
@@ -1,4 +1,5 @@
+from ._tasks import IconTask, ShellTask
 from .graph_items import Cycle, Data, GraphItem, Task
 from .workflow import Workflow
 
-__all__ = ["Workflow", "GraphItem", "Data", "Task", "Cycle"]
+__all__ = ["Workflow", "GraphItem", "Data", "Task", "Cycle", "ShellTask", "IconTask"]
diff --git a/src/sirocco/core/_tasks/__init__.py b/src/sirocco/core/_tasks/__init__.py
index ee398448..d26e4772 100644
--- a/src/sirocco/core/_tasks/__init__.py
+++ b/src/sirocco/core/_tasks/__init__.py
@@ -1,3 +1,4 @@
-from . import icon_task, shell_task
+from .icon_task import IconTask
+from .shell_task import ShellTask
 
-__all__ = ["icon_task", "shell_task"]
+__all__ = ["IconTask", "ShellTask"]
diff --git a/src/sirocco/core/graph_items.py b/src/sirocco/core/graph_items.py
index 371089d9..4ceba492 100644
--- a/src/sirocco/core/graph_items.py
+++ b/src/sirocco/core/graph_items.py
@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass, field
 from itertools import chain, product
-from typing import TYPE_CHECKING, Any, ClassVar, Self, TypeAlias
+from typing import TYPE_CHECKING, Any, ClassVar, Self, TypeAlias, TypeVar, cast
 
 from sirocco.parsing._yaml_data_models import (
     ConfigAvailableData,
@@ -15,9 +15,12 @@
     from datetime import datetime
     from pathlib import Path
 
+    from termcolor._types import Color
+
     from sirocco.parsing._yaml_data_models import (
         ConfigBaseData,
         ConfigCycleTask,
+        ConfigCycleTaskWaitOn,
         ConfigTask,
         TargetNodesBaseModel,
     )
@@ -27,17 +30,20 @@
 class GraphItem:
     """base class for Data Tasks and Cycles"""
 
-    color: ClassVar[str]
+    color: ClassVar[Color]
 
     name: str
     coordinates: dict
 
 
+GRAPH_ITEM_T = TypeVar("GRAPH_ITEM_T", bound=GraphItem)
+
+
 @dataclass(kw_only=True)
 class Data(ConfigBaseDataSpecs, GraphItem):
     """Internal representation of a data node"""
 
-    color: ClassVar[str] = field(default="light_blue", repr=False)
+    color: ClassVar[Color] = field(default="light_blue", repr=False)
 
     available: bool
 
@@ -61,15 +67,17 @@ class Task(ConfigBaseTaskSpecs, GraphItem):
     """Internal representation of a task node"""
 
     plugin_classes: ClassVar[dict[str, type]] = field(default={}, repr=False)
-    color: ClassVar[str] = field(default="light_red", repr=False)
+    color: ClassVar[Color] = field(default="light_red", repr=False)
 
     inputs: list[BoundData] = field(default_factory=list)
     outputs: list[Data] = field(default_factory=list)
     wait_on: list[Task] = field(default_factory=list)
-    config_rootdir: Path | None = None
+    config_rootdir: Path
     start_date: datetime | None = None
     end_date: datetime | None = None
 
+    _wait_on_specs: list[ConfigCycleTaskWaitOn] = field(default_factory=list, repr=False)
+
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
         if cls.plugin in Task.plugin_classes:
@@ -118,7 +126,7 @@ def from_config(
 
         return new
 
-    def link_wait_on_tasks(self, taskstore: Store):
+    def link_wait_on_tasks(self, taskstore: Store[Task]) -> None:
         self.wait_on = list(
             chain(
                 *(
@@ -133,24 +141,24 @@ def link_wait_on_tasks(self, taskstore: Store):
 class Cycle(GraphItem):
     """Internal reprenstation of a cycle"""
 
-    color: ClassVar[str] = field(default="light_green", repr=False)
+    color: ClassVar[Color] = field(default="light_green", repr=False)
 
     tasks: list[Task]
 
 
-class Array:
-    """Dictionnary of GraphItem objects accessed by arbitrary dimensions"""
+class Array[GRAPH_ITEM_T]:
+    """Dictionnary of GRAPH_ITEM_T objects accessed by arbitrary dimensions"""
 
     def __init__(self, name: str) -> None:
         self._name = name
-        self._dims: tuple[str] | None = None
-        self._axes: dict | None = None
-        self._dict: dict[tuple, GraphItem] | None = None
+        self._dims: tuple[str, ...] = ()
+        self._axes: dict[str, set] = {}
+        self._dict: dict[tuple, GRAPH_ITEM_T] = {}
 
-    def __setitem__(self, coordinates: dict, value: GraphItem) -> None:
+    def __setitem__(self, coordinates: dict, value: GRAPH_ITEM_T) -> None:
         # First access: set axes and initialize dictionnary
         input_dims = tuple(coordinates.keys())
-        if self._dims is None:
+        if self._dims == ():
             self._dims = input_dims
             self._axes = {k: set() for k in self._dims}
             self._dict = {}
@@ -171,7 +179,7 @@ def __setitem__(self, coordinates: dict, value: GraphItem) -> None:
         # Set item
         self._dict[key] = value
 
-    def __getitem__(self, coordinates: dict) -> GraphItem:
+    def __getitem__(self, coordinates: dict) -> GRAPH_ITEM_T:
         if self._dims != (input_dims := tuple(coordinates.keys())):
             msg = f"Array {self._name}: coordinate names {input_dims} don't match Array dimensions {self._dims}"
             raise KeyError(msg)
@@ -179,7 +187,7 @@ def __getitem__(self, coordinates: dict) -> GraphItem:
         key = tuple(coordinates[dim] for dim in self._dims)
         return self._dict[key]
 
-    def iter_from_cycle_spec(self, spec: TargetNodesBaseModel, reference: dict) -> Iterator[GraphItem]:
+    def iter_from_cycle_spec(self, spec: TargetNodesBaseModel, reference: dict) -> Iterator[GRAPH_ITEM_T]:
         # Check date references
         if "date" not in self._dims and (spec.lag or spec.date):
             msg = f"Array {self._name} has no date dimension, cannot be referenced by dates"
@@ -205,26 +213,24 @@ def _resolve_target_dim(self, spec: TargetNodesBaseModel, dim: str, reference: A
         else:
             yield from self._axes[dim]
 
-    def __iter__(self) -> Iterator[GraphItem]:
+    def __iter__(self) -> Iterator[GRAPH_ITEM_T]:
         yield from self._dict.values()
 
 
-class Store:
-    """Container for GraphItem Arrays"""
+class Store[GRAPH_ITEM_T]:
+    """Container for GRAPH_ITEM_T Arrays"""
 
-    def __init__(self):
-        self._dict: dict[str, Array] = {}
+    def __init__(self) -> None:
+        self._dict: dict[str, Array[GRAPH_ITEM_T]] = {}
 
-    def add(self, item) -> None:
-        if not isinstance(item, GraphItem):
-            msg = "items in a Store must be of instance GraphItem"
-            raise TypeError(msg)
-        name, coordinates = item.name, item.coordinates
+    def add(self, item: GRAPH_ITEM_T) -> None:
+        graph_item = cast(GraphItem, item)  # mypy can somehow not deduce this
+        name, coordinates = graph_item.name, graph_item.coordinates
         if name not in self._dict:
-            self._dict[name] = Array(name)
+            self._dict[name] = Array[GRAPH_ITEM_T](name)
         self._dict[name][coordinates] = item
 
-    def __getitem__(self, key: tuple[str, dict]) -> GraphItem:
+    def __getitem__(self, key: tuple[str, dict]) -> GRAPH_ITEM_T:
         name, coordinates = key
         if "date" in coordinates and coordinates["date"] is None:
             del coordinates["date"]
@@ -233,7 +239,7 @@ def __getitem__(self, key: tuple[str, dict]) -> GraphItem:
             raise KeyError(msg)
         return self._dict[name][coordinates]
 
-    def iter_from_cycle_spec(self, spec: TargetNodesBaseModel, reference: dict) -> Iterator[GraphItem]:
+    def iter_from_cycle_spec(self, spec: TargetNodesBaseModel, reference: dict) -> Iterator[GRAPH_ITEM_T]:
         # Check if target items should be querried at all
         if (when := spec.when) is not None:
             if (ref_date := reference.get("date")) is None:
@@ -248,5 +254,5 @@ def iter_from_cycle_spec(self, spec: TargetNodesBaseModel, reference: dict) -> I
         # Yield items
         yield from self._dict[spec.name].iter_from_cycle_spec(spec, reference)
 
-    def __iter__(self) -> Iterator[GraphItem]:
+    def __iter__(self) -> Iterator[GRAPH_ITEM_T]:
         yield from chain(*(self._dict.values()))
diff --git a/src/sirocco/core/workflow.py b/src/sirocco/core/workflow.py
index c394ad4a..dd6c0e44 100644
--- a/src/sirocco/core/workflow.py
+++ b/src/sirocco/core/workflow.py
@@ -5,6 +5,7 @@
 
 from sirocco.core.graph_items import Cycle, Data, Store, Task
 from sirocco.parsing._yaml_data_models import (
+    ConfigBaseData,
     ConfigWorkflow,
 )
 
@@ -14,10 +15,8 @@
     from pathlib import Path
 
     from sirocco.parsing._yaml_data_models import (
-        ConfigAvailableData,
         ConfigCycle,
         ConfigData,
-        ConfigGeneratedData,
         ConfigTask,
     )
 
@@ -37,13 +36,11 @@ def __init__(
         self.name: str = name
         self.config_rootdir: Path = config_rootdir
 
-        self.tasks: Store = Store()
-        self.data: Store = Store()
-        self.cycles: Store = Store()
+        self.tasks: Store[Task] = Store()
+        self.data: Store[Data] = Store()
+        self.cycles: Store[Cycle] = Store()
 
-        data_dict: dict[str, ConfigAvailableData | ConfigGeneratedData] = {
-            data.name: data for data in chain(data.available, data.generated)
-        }
+        data_dict: dict[str, ConfigBaseData] = {data.name: data for data in chain(data.available, data.generated)}
         task_dict: dict[str, ConfigTask] = {task.name: task for task in tasks}
 
         # Function to iterate over date and parameter combinations
@@ -52,9 +49,9 @@ def iter_coordinates(param_refs: list, date: datetime | None = None) -> Iterator
             yield from (dict(zip(space.keys(), x, strict=False)) for x in product(*space.values()))
 
         # 1 - create availalbe data nodes
-        for data_config in data.available:
-            for coordinates in iter_coordinates(param_refs=data_config.parameters, date=None):
-                self.data.add(Data.from_config(config=data_config, coordinates=coordinates))
+        for available_data_config in data.available:
+            for coordinates in iter_coordinates(param_refs=available_data_config.parameters, date=None):
+                self.data.add(Data.from_config(config=available_data_config, coordinates=coordinates))
 
         # 2 - create output data nodes
         for cycle_config in cycles:
@@ -100,9 +97,9 @@ def iter_coordinates(param_refs: list, date: datetime | None = None) -> Iterator
             task.link_wait_on_tasks(self.tasks)
 
     @staticmethod
-    def cycle_dates(cycle_config: ConfigCycle) -> Iterator[datetime]:
+    def cycle_dates(cycle_config: ConfigCycle) -> Iterator[datetime | None]:
         yield (date := cycle_config.start_date)
-        if cycle_config.period is not None:
+        if cycle_config.period is not None and date is not None and cycle_config.end_date is not None:
             while (date := date + cycle_config.period) < cycle_config.end_date:
                 yield date
 
diff --git a/src/sirocco/parsing/_yaml_data_models.py b/src/sirocco/parsing/_yaml_data_models.py
index c84bee67..657360ca 100644
--- a/src/sirocco/parsing/_yaml_data_models.py
+++ b/src/sirocco/parsing/_yaml_data_models.py
@@ -122,9 +122,9 @@ def check_before_after_at_combination(cls, data: Any) -> Any:
 
     @field_validator("before", "after", "at", mode="before")
     @classmethod
-    def convert_datetime(cls, value) -> datetime:
-        if value is None:
-            return None
+    def convert_datetime(cls, value: datetime | str | None) -> datetime | None:
+        if value is None or isinstance(value, datetime):
+            return value
         return datetime.fromisoformat(value)
 
 
@@ -192,53 +192,46 @@ class ConfigCycleTaskOutput(_NamedBaseModel):
     """
 
 
-class ConfigCycleTask(_NamedBaseModel):
-    """
-    To create an instance of a task in a cycle defined in a workflow file.
-    """
+NAMED_BASE_T = typing.TypeVar("NAMED_BASE_T", bound=_NamedBaseModel)
 
-    inputs: list[ConfigCycleTaskInput | str] | None = Field(default_factory=list)
-    outputs: list[ConfigCycleTaskOutput | str] | None = Field(default_factory=list)
-    wait_on: list[ConfigCycleTaskWaitOn | str] | None = Field(default_factory=list)
 
-    @field_validator("inputs", mode="before")
-    @classmethod
-    def convert_cycle_task_inputs(cls, values) -> list[ConfigCycleTaskInput]:
-        inputs = []
+def make_named_model_list_converter(
+    cls: type[NAMED_BASE_T],
+) -> typing.Callable[[list[NAMED_BASE_T | str | dict] | None], list[NAMED_BASE_T]]:
+    def convert_named_model_list(values: list[NAMED_BASE_T | str | dict] | None) -> list[NAMED_BASE_T]:
+        inputs: list[NAMED_BASE_T] = []
         if values is None:
             return inputs
         for value in values:
-            if isinstance(value, str):
-                inputs.append({value: None})
-            elif isinstance(value, dict):
-                inputs.append(value)
+            match value:
+                case str():
+                    inputs.append(cls(name=value))
+                case dict():
+                    inputs.append(cls(**value))
+                case _NamedBaseModel():
+                    inputs.append(value)
+                case _:
+                    msg = "Unsupported Type"
+                    raise TypeError(msg)
         return inputs
 
-    @field_validator("outputs", mode="before")
-    @classmethod
-    def convert_cycle_task_outputs(cls, values) -> list[ConfigCycleTaskOutput]:
-        outputs = []
-        if values is None:
-            return outputs
-        for value in values:
-            if isinstance(value, str):
-                outputs.append({value: None})
-            elif isinstance(value, dict):
-                outputs.append(value)
-        return outputs
+    return convert_named_model_list
 
-    @field_validator("wait_on", mode="before")
-    @classmethod
-    def convert_cycle_task_wait_on(cls, values) -> list[ConfigCycleTaskWaitOn]:
-        wait_on = []
-        if values is None:
-            return wait_on
-        for value in values:
-            if isinstance(value, str):
-                wait_on.append({value: None})
-            elif isinstance(value, dict):
-                wait_on.append(value)
-        return wait_on
+
+class ConfigCycleTask(_NamedBaseModel):
+    """
+    To create an instance of a task in a cycle defined in a workflow file.
+    """
+
+    inputs: Annotated[
+        list[ConfigCycleTaskInput], BeforeValidator(make_named_model_list_converter(ConfigCycleTaskInput))
+    ] = []
+    outputs: Annotated[
+        list[ConfigCycleTaskOutput], BeforeValidator(make_named_model_list_converter(ConfigCycleTaskOutput))
+    ] = []
+    wait_on: Annotated[
+        list[ConfigCycleTaskWaitOn], BeforeValidator(make_named_model_list_converter(ConfigCycleTaskWaitOn))
+    ] = []
 
 
 class ConfigCycle(_NamedBaseModel):
diff --git a/src/sirocco/pretty_print.py b/src/sirocco/pretty_print.py
index 7e219573..1dfa1ad8 100644
--- a/src/sirocco/pretty_print.py
+++ b/src/sirocco/pretty_print.py
@@ -85,10 +85,12 @@ def format_basic(self, obj: core.GraphItem) -> str:
         Example:
 
         >>> from datetime import datetime
+        >>> import pathlib
         >>> print(
         ...     PrettyPrinter().format_basic(
         ...         core.Task(
         ...             name="foo",
+        ...             config_rootdir=pathlib.Path("."),
         ...             coordinates={"date": datetime(1000, 1, 1).date()},
         ...         )
         ...     )
diff --git a/src/sirocco/py.typed b/src/sirocco/py.typed
new file mode 100644
index 00000000..e69de29b
diff --git a/src/sirocco/vizgraph.py b/src/sirocco/vizgraph.py
index 1ccd3cfe..6395cc23 100644
--- a/src/sirocco/vizgraph.py
+++ b/src/sirocco/vizgraph.py
@@ -9,7 +9,7 @@
 from pygraphviz import AGraph
 
 if TYPE_CHECKING:
-    from sirocco.core import Store
+    from sirocco.core.graph_items import Store
 from sirocco.core import Workflow
 
 
@@ -18,7 +18,7 @@ def hsv_to_hex(h: float, s: float, v: float) -> str:
     return "#{:02x}{:02x}{:02x}".format(*map(round, (255 * r, 255 * g, 255 * b)))
 
 
-def node_colors(h: float) -> dict[str:str]:
+def node_colors(h: float) -> dict[str, str]:
     fill = hsv_to_hex(h / 365, 0.15, 1)
     border = hsv_to_hex(h / 365, 1, 0.20)
     font = hsv_to_hex(h / 365, 1, 0.15)
@@ -28,16 +28,16 @@ def node_colors(h: float) -> dict[str:str]:
 class VizGraph:
     """Class for visualizing a Sirocco workflow"""
 
-    node_base_kw: ClassVar[dict[str:Any]] = {"style": "filled", "fontname": "Fira Sans", "fontsize": 14, "penwidth": 2}
-    edge_base_kw: ClassVar[dict[str:Any]] = {"color": "#77767B", "penwidth": 1.5}
-    data_node_base_kw: ClassVar[dict[str:Any]] = node_base_kw | {"shape": "ellipse"}
+    node_base_kw: ClassVar[dict[str, Any]] = {"style": "filled", "fontname": "Fira Sans", "fontsize": 14, "penwidth": 2}
+    edge_base_kw: ClassVar[dict[str, Any]] = {"color": "#77767B", "penwidth": 1.5}
+    data_node_base_kw: ClassVar[dict[str, Any]] = node_base_kw | {"shape": "ellipse"}
 
-    data_av_node_kw: ClassVar[dict[str:Any]] = data_node_base_kw | node_colors(116)
-    data_gen_node_kw: ClassVar[dict[str:Any]] = data_node_base_kw | node_colors(214)
-    task_node_kw: ClassVar[dict[str:Any]] = node_base_kw | {"shape": "box"} | node_colors(354)
-    io_edge_kw: ClassVar[dict[str:Any]] = edge_base_kw
-    wait_on_edge_kw: ClassVar[dict[str:Any]] = edge_base_kw | {"style": "dashed"}
-    cluster_kw: ClassVar[dict[str:Any]] = {"bgcolor": "#F6F5F4", "color": None, "fontsize": 16}
+    data_av_node_kw: ClassVar[dict[str, Any]] = data_node_base_kw | node_colors(116)
+    data_gen_node_kw: ClassVar[dict[str, Any]] = data_node_base_kw | node_colors(214)
+    task_node_kw: ClassVar[dict[str, Any]] = node_base_kw | {"shape": "box"} | node_colors(354)
+    io_edge_kw: ClassVar[dict[str, Any]] = edge_base_kw
+    wait_on_edge_kw: ClassVar[dict[str, Any]] = edge_base_kw | {"style": "dashed"}
+    cluster_kw: ClassVar[dict[str, Any]] = {"bgcolor": "#F6F5F4", "color": None, "fontsize": 16}
 
     def __init__(self, name: str, cycles: Store, data: Store) -> None:
         self.name = name
diff --git a/src/sirocco/workgraph.py b/src/sirocco/workgraph.py
index 2e44cde7..ea360639 100644
--- a/src/sirocco/workgraph.py
+++ b/src/sirocco/workgraph.py
@@ -9,15 +9,11 @@
 from aiida.common.exceptions import NotExistent
 from aiida_workgraph import WorkGraph
 
-from sirocco.core._tasks.icon_task import IconTask
-from sirocco.core._tasks.shell_task import ShellTask
+from sirocco import core
 
 if TYPE_CHECKING:
     from aiida_workgraph.socket import TaskSocket  # type: ignore[import-untyped]
 
-    from sirocco import core
-    from sirocco.core import graph_items
-
 
 # This is a workaround required when splitting the initialization of the task and its linked nodes Merging this into
 # aiida-workgraph properly would require significant changes see issues
@@ -126,7 +122,7 @@ def replace_invalid_chars_in_label(label: str) -> str:
         return label
 
     @staticmethod
-    def get_aiida_label_from_graph_item(obj: graph_items.GraphItem) -> str:
+    def get_aiida_label_from_graph_item(obj: core.GraphItem) -> str:
         """Returns a unique AiiDA label for the given graph item.
 
         The graph item object is uniquely determined by its name and its coordinates. There is the possibility that
@@ -136,7 +132,7 @@ def get_aiida_label_from_graph_item(obj: graph_items.GraphItem) -> str:
             f"{obj.name}" + "__".join(f"_{key}_{value}" for key, value in obj.coordinates.items())
         )
 
-    def _add_aiida_input_data_node(self, data: graph_items.Data):
+    def _add_aiida_input_data_node(self, data: core.Data):
         """
         Create an `aiida.orm.Data` instance from the provided graph item.
         """
@@ -179,15 +175,15 @@ def _add_tasks(self):
                 self._link_input_nodes_to_task(task, input_)
             self._link_arguments_to_task(task)
 
-    def _create_task_node(self, task: graph_items.Task):
+    def _create_task_node(self, task: core.Task):
         label = AiidaWorkGraph.get_aiida_label_from_graph_item(task)
-        if isinstance(task, ShellTask):
+        if isinstance(task, core.ShellTask):
             command_path = Path(task.command)
             command_full_path = task.command if command_path.is_absolute() else task.config_rootdir / command_path
             command = str(command_full_path)
 
             # metadata
-            metadata = {}
+            metadata: dict[str, Any] = {}
             ## Source file
             env_source_paths = [
                 env_source_path
@@ -220,14 +216,14 @@ def _create_task_node(self, task: graph_items.Task):
 
             self._aiida_task_nodes[label] = workgraph_task
 
-        elif isinstance(task, IconTask):
+        elif isinstance(task, core.IconTask):
             exc = "IconTask not implemented yet."
             raise NotImplementedError(exc)
         else:
             exc = f"Task: {task.name} not implemented yet."
             raise NotImplementedError(exc)
 
-    def _link_wait_on_to_task(self, task: graph_items.Task):
+    def _link_wait_on_to_task(self, task: core.Task):
         label = AiidaWorkGraph.get_aiida_label_from_graph_item(task)
         workgraph_task = self._aiida_task_nodes[label]
         wait_on_tasks = []
@@ -236,7 +232,7 @@ def _link_wait_on_to_task(self, task: graph_items.Task):
             wait_on_tasks.append(self._aiida_task_nodes[wait_on_task_label])
         workgraph_task.wait = wait_on_tasks
 
-    def _link_input_nodes_to_task(self, task: graph_items.Task, input_: graph_items.Data):
+    def _link_input_nodes_to_task(self, task: core.Task, input_: core.Data):
         """Links the input to the workgraph task."""
         task_label = AiidaWorkGraph.get_aiida_label_from_graph_item(task)
         input_label = AiidaWorkGraph.get_aiida_label_from_graph_item(input_)
@@ -259,7 +255,7 @@ def _link_input_nodes_to_task(self, task: graph_items.Task, input_: graph_items.
             )
             raise ValueError(msg)
 
-    def _link_arguments_to_task(self, task: graph_items.Task):
+    def _link_arguments_to_task(self, task: core.Task):
         """Links the arguments to the workgraph task.
 
         Parses `cli_arguments` of the graph item task and links all arguments to the task node. It only adds arguments
@@ -277,6 +273,8 @@ def _link_arguments_to_task(self, task: graph_items.Task):
         name_to_input_map = {input_.name: input_ for input_, _ in task.inputs}
         # we track the linked input arguments, to ensure that all linked input nodes got linked arguments
         linked_input_args = []
+        if not isinstance(task, core.ShellTask):
+            raise TypeError
         for arg in task.cli_arguments:
             if arg.references_data_item:
                 # We only add an input argument to the args if it has been added to the nodes
@@ -298,7 +296,7 @@ def _link_arguments_to_task(self, task: graph_items.Task):
                 input_label = AiidaWorkGraph.get_aiida_label_from_graph_item(input_)
                 workgraph_task_arguments.value.append(f"{{{input_label}}}")
 
-    def _link_output_nodes_to_task(self, task: graph_items.Task, output: graph_items.Data):
+    def _link_output_nodes_to_task(self, task: core.Task, output: core.Data):
         """Links the output to the workgraph task."""
 
         workgraph_task = self._aiida_task_nodes[AiidaWorkGraph.get_aiida_label_from_graph_item(task)]
diff --git a/tests/conftest.py b/tests/conftest.py
index 97184f56..bf0448c7 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -16,8 +16,8 @@ def minimal_config() -> models.ConfigWorkflow:
     return models.ConfigWorkflow(
         name="minimal",
         rootdir=pathlib.Path("minimal"),
-        cycles=[models.ConfigCycle(minimal={"tasks": [models.ConfigCycleTask(some_task={})]})],
-        tasks=[models.ConfigShellTask(some_task={"plugin": "shell"})],
+        cycles=[models.ConfigCycle(name="minimal", tasks=[models.ConfigCycleTask(name="some_task")])],
+        tasks=[models.ConfigShellTask(name="some_task")],
         data=models.ConfigData(
             available=[models.ConfigAvailableData(name="foo", type=models.DataType.FILE, src="foo.txt")],
             generated=[models.ConfigGeneratedData(name="bar", type=models.DataType.DIR, src="bar")],