Improve yaml models: ConfigTasks

src/sirocco/core/_tasks/icon_task.py

@@ -3,15 +3,19 @@
 import re
 from dataclasses import dataclass, field
 from pathlib import Path
+from typing import TYPE_CHECKING, Any, Self
 
 import f90nml
 
 from sirocco.core.graph_items import Task
-from sirocco.parsing._yaml_data_models import ConfigIconTaskSpecs
+from sirocco.parsing import yaml_data_models as models
+
+if TYPE_CHECKING:
+    from sirocco.parsing.yaml_data_models import ConfigTask
 
 
 @dataclass(kw_only=True)
-class IconTask(ConfigIconTaskSpecs, Task):
+class IconTask(models.ConfigIconTaskSpecs, Task):
     core_namelists: dict[str, f90nml.Namelist] = field(default_factory=dict)
 
     def init_core_namelists(self):
@@ -93,3 +97,20 @@ def section_index(section_name) -> tuple[str, int | None]:
         if m := multi_section_pattern.match(section_name):
             return m.group(1), int(m.group(2)) - 1
         return section_name, None
+
+    @classmethod
+    def build_from_config(cls: type[Self], config: ConfigTask, **kwargs: Any) -> Self:
+        config_kwargs = dict(config)
+        del config_kwargs["parameters"]
+        # The following check is here for type checkers.
+        # We don't want to narrow the type in the signature, as that would break liskov substitution.
+        # We guarantee elsewhere this is called with the correct type at runtime
+        if not isinstance(config, models.ConfigIconTask):
+            raise TypeError
+        config_kwargs["namelists"] = {
+            nml.path.name: models.NamelistSpec(**nml.model_dump()) for nml in config.namelists
+        }
+        return cls(
+            **kwargs,
+            **config_kwargs,
+        )

src/sirocco/core/_tasks/shell_task.py

@@ -3,9 +3,8 @@
 from dataclasses import dataclass
 
 from sirocco.core.graph_items import Task
-from sirocco.parsing._yaml_data_models import ConfigShellTaskSpecs
+from sirocco.parsing.yaml_data_models import ConfigShellTaskSpecs
 
 
 @dataclass(kw_only=True)
-class ShellTask(ConfigShellTaskSpecs, Task):
-    pass
+class ShellTask(ConfigShellTaskSpecs, Task): ...

src/sirocco/core/graph_items.py

@@ -4,7 +4,7 @@
 from itertools import chain, product
 from typing import TYPE_CHECKING, Any, ClassVar, Self, TypeAlias, TypeVar, cast
 
-from sirocco.parsing._yaml_data_models import (
+from sirocco.parsing.yaml_data_models import (
     ConfigAvailableData,
     ConfigBaseDataSpecs,
     ConfigBaseTaskSpecs,
@@ -17,7 +17,7 @@
 
     from termcolor._types import Color
 
-    from sirocco.parsing._yaml_data_models import (
+    from sirocco.parsing.yaml_data_models import (
         ConfigBaseData,
         ConfigCycleTask,
         ConfigCycleTaskWaitOn,
@@ -66,7 +66,7 @@ def from_config(cls, config: ConfigBaseData, coordinates: dict) -> Self:
 class Task(ConfigBaseTaskSpecs, GraphItem):
     """Internal representation of a task node"""
 
-    plugin_classes: ClassVar[dict[str, type]] = field(default={}, repr=False)
+    plugin_classes: ClassVar[dict[str, type[Self]]] = field(default={}, repr=False)
     color: ClassVar[Color] = field(default="light_red", repr=False)
 
     inputs: list[BoundData] = field(default_factory=list)
@@ -87,7 +87,7 @@ def __init_subclass__(cls, **kwargs):
 
     @classmethod
     def from_config(
-        cls,
+        cls: type[Self],
         config: ConfigTask,
         config_rootdir: Path,
         start_date: datetime | None,
@@ -102,22 +102,19 @@ def from_config(
             for data_node in datastore.iter_from_cycle_spec(input_spec, coordinates)
         ]
         outputs = [datastore[output_spec.name, coordinates] for output_spec in graph_spec.outputs]
-        # use the fact that pydantic models can be turned into dicts easily
-        cls_config = dict(config)
-        del cls_config["parameters"]
         if (plugin_cls := Task.plugin_classes.get(type(config).plugin, None)) is None:
             msg = f"Plugin {type(config).plugin!r} is not supported."
             raise ValueError(msg)
 
-        new = plugin_cls(
+        new = plugin_cls.build_from_config(
+            config,
             config_rootdir=config_rootdir,
             coordinates=coordinates,
             start_date=start_date,
             end_date=end_date,
             inputs=inputs,
             outputs=outputs,
-            **cls_config,
-        )  # this works because dataclass has generated this init for us
+        )
 
         # Store for actual linking in link_wait_on_tasks() once all tasks are created
         new._wait_on_specs = graph_spec.wait_on  # noqa: SLF001 we don't have access to self in a dataclass
@@ -126,6 +123,12 @@ def from_config(
 
         return new
 
+    @classmethod
+    def build_from_config(cls: type[Self], config: ConfigTask, **kwargs: Any) -> Self:
+        config_kwargs = dict(config)
+        del config_kwargs["parameters"]
+        return cls(**kwargs, **config_kwargs)
+
     def link_wait_on_tasks(self, taskstore: Store[Task]) -> None:
         self.wait_on = list(
             chain(

src/sirocco/core/workflow.py

@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, Self
 
 from sirocco.core.graph_items import Cycle, Data, Store, Task
-from sirocco.parsing._yaml_data_models import (
+from sirocco.parsing.yaml_data_models import (
     ConfigBaseData,
     ConfigWorkflow,
 )
@@ -14,7 +14,7 @@
     from datetime import datetime
     from pathlib import Path
 
-    from sirocco.parsing._yaml_data_models import (
+    from sirocco.parsing.yaml_data_models import (
         ConfigCycle,
         ConfigData,
         ConfigTask,

src/sirocco/parsing/__init__.py

@@ -1,4 +1,4 @@
-from ._yaml_data_models import (
+from .yaml_data_models import (
     ConfigWorkflow,
 )