Restore duration plugin and ty config

Author: tobiasraabe

justfile

@@ -12,7 +12,7 @@ test-cov *FLAGS:
 
 # Run type checking
 typing:
-    uv run --group typing --group test --isolated ty check src/ tests/
+    uv run --group typing --group test --isolated ty check
 
 # Run linting
 lint:

pyproject.toml

@@ -179,6 +179,7 @@ include = [
 unused-ignore-comment = "ignore"
 
 [tool.ty.src]
+include = ["src", "tests"]
 exclude = ["src/_pytask/_hashlib.py"]
 
 [tool.ty.terminal]

src/_pytask/profile.py

@@ -54,7 +54,9 @@ def pytask_extend_command_line_interface(cli: click.Group) -> None:
 @hookimpl
 def pytask_post_parse(config: dict[str, Any]) -> None:
     """Register the export option."""
-    config["pm"].register(ProfilePlugin(RuntimeState.from_root(config["root"])))
+    runtime_state = RuntimeState.from_root(config["root"])
+    config["pm"].register(ProfilePlugin(runtime_state))
+    config["pm"].register(DurationNameSpace(runtime_state))
     config["pm"].register(ExportNameSpace)
     config["pm"].register(FileSizeNameSpace)
 
@@ -86,15 +88,6 @@ def pytask_execute_task_process_report(
         if report.outcome == TaskOutcome.SUCCESS and duration is not None:
             self.runtime_state.update_task(task, *duration)
 
-    @hookimpl
-    def pytask_profile_add_info_on_task(
-        self, session: Session, tasks: list[PTask], profile: dict[str, dict[str, Any]]
-    ) -> None:
-        """Add the runtime for tasks to the profile."""
-        _ = session
-        for name, duration in self._collect_runtimes(tasks).items():
-            profile[name]["Duration (in s)"] = round(duration, 2)
-
     @hookimpl
     def pytask_unconfigure(self, session: Session) -> None:
         """Flush runtime information on normal build exits."""
@@ -104,13 +97,23 @@ def pytask_unconfigure(self, session: Session) -> None:
             return
         self.runtime_state.flush()
 
-    def _collect_runtimes(self, tasks: list[PTask]) -> dict[str, float]:
-        """Collect runtimes."""
-        return {
-            task.name: duration
-            for task in tasks
-            if (duration := self.runtime_state.get_duration(task)) is not None
-        }
+
+class DurationNameSpace:
+    """A namespace for adding durations to the profile."""
+
+    def __init__(self, runtime_state: RuntimeState) -> None:
+        self.runtime_state = runtime_state
+
+    @hookimpl
+    def pytask_profile_add_info_on_task(
+        self, session: Session, tasks: list[PTask], profile: dict[str, dict[str, Any]]
+    ) -> None:
+        """Add the runtime for tasks to the profile."""
+        _ = session
+        for task in tasks:
+            duration = self.runtime_state.get_duration(task)
+            if duration is not None:
+                profile[task.name]["Duration (in s)"] = round(duration, 2)
 
 
 @click.command(cls=ColoredCommand)

tests/test_runtime_store.py

@@ -1,15 +1,37 @@
 from __future__ import annotations
 
-from types import SimpleNamespace
+from typing import Any
 
 import pytest
 
 from _pytask.runtime_store import RuntimeState
 
 
+class DummyTask:
+    def __init__(self, name: str) -> None:
+        self.name = name
+        self.depends_on = {}
+        self.produces = {}
+        self.function = lambda: None
+        self.markers = []
+        self.report_sections = []
+        self.attributes = {}
+
+    @property
+    def signature(self) -> str:
+        return self.name
+
+    def state(self) -> str | None:
+        return None
+
+    def execute(self, **kwargs: Any) -> Any:
+        _ = kwargs
+        return None
+
+
 def test_runtime_state_recovers_from_journal(tmp_path):
     tmp_path.joinpath(".pytask").mkdir()
-    task = SimpleNamespace(name="task_example")
+    task = DummyTask(name="task_example")
 
     state = RuntimeState.from_root(tmp_path)
     state.update_task(task, 1.0, 3.0)
@@ -20,7 +42,7 @@ def test_runtime_state_recovers_from_journal(tmp_path):
 
 def test_runtime_state_flushes_journal(tmp_path):
     tmp_path.joinpath(".pytask").mkdir()
-    task = SimpleNamespace(name="task_example")
+    task = DummyTask(name="task_example")
 
     state = RuntimeState.from_root(tmp_path)
     state.update_task(task, 2.0, 5.5)