Add TaskStage API support with models, resource, tests and examples

Author: TanyaSingh369-svg

examples/task_stage_example.py

@@ -0,0 +1,38 @@
+"""
+Example usage of TaskStages API
+
+This demonstrates how to:
+- Read a task stage
+- List task stages for a run
+- Override a task stage
+"""
+
+from pytfe.client import TFEClient
+
+# Initialize client (make sure your auth/env is configured)
+client = TFEClient()
+
+# ---------------------------
+# Read a task stage
+# ---------------------------
+# Fetch a single task stage by ID
+# Replace "ts-123" with a real task stage ID
+stage = client.task_stages.read("ts-abc123xyz")
+print(stage)
+
+
+# ---------------------------
+# List task stages for a run
+# ---------------------------
+# Fetch all task stages for a run
+# Replace "run-123" with a real run ID
+# for stage in client.task_stages.list("run-123"):
+#     print(stage)
+
+
+# ---------------------------
+# Override a task stage
+# ---------------------------
+# Override a task stage (if allowed)
+# Replace "ts-123" with a real task stage ID
+# client.task_stages.override("ts-123", comment="Approved")

src/pytfe/client.py

@@ -34,6 +34,7 @@
 from .resources.run_event import RunEvents
 from .resources.run_task import RunTasks
 from .resources.run_trigger import RunTriggers
+from .resources.task_stage import TaskStages
 from .resources.ssh_keys import SSHKeys
 from .resources.state_version_outputs import StateVersionOutputs
 from .resources.state_versions import StateVersions
@@ -94,6 +95,7 @@ def __init__(self, config: TFEConfig | None = None):
         self.run_tasks = RunTasks(self._transport)
         self.run_triggers = RunTriggers(self._transport)
         self.runs = Runs(self._transport)
+        self.task_stages = TaskStages(self._transport)
         self.query_runs = QueryRuns(self._transport)
         self.run_events = RunEvents(self._transport)
         self.policies = Policies(self._transport)

src/pytfe/models/task_result.py

@@ -0,0 +1,60 @@
+# Copyright IBM Corp. 2025, 2026
+# SPDX-License-Identifier: MPL-2.0
+
+from __future__ import annotations
+
+from datetime import datetime
+from enum import Enum
+from typing import Optional, TYPE_CHECKING
+
+from pydantic import BaseModel, ConfigDict, Field
+
+class TaskResultStatus(str, Enum):
+    passed = "passed"
+    failed = "failed"
+    pending = "pending"
+    running = "running"
+    unreachable = "unreachable"
+    errored = "errored"
+
+
+class TaskEnforcementLevel(str, Enum):
+    advisory = "advisory"
+    mandatory = "mandatory"
+
+
+class TaskResultStatusTimestamps(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
+
+    errored_at: Optional[datetime] = Field(None, alias="errored-at")
+    running_at: Optional[datetime] = Field(None, alias="running-at")
+    canceled_at: Optional[datetime] = Field(None, alias="canceled-at")
+    failed_at: Optional[datetime] = Field(None, alias="failed-at")
+    passed_at: Optional[datetime] = Field(None, alias="passed-at")
+
+
+class TaskResult(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
+
+    id: str
+    status: TaskResultStatus = Field(..., alias="status")
+    message: str = Field(..., alias="message")
+
+    status_timestamps: TaskResultStatusTimestamps = Field(..., alias="status-timestamps")
+
+    url: str = Field(..., alias="url")
+
+    created_at: datetime = Field(..., alias="created-at")
+    updated_at: datetime = Field(..., alias="updated-at")
+
+    task_id: str = Field(..., alias="task-id")
+    task_name: str = Field(..., alias="task-name")
+    task_url: str = Field(..., alias="task-url")
+
+    workspace_task_id: str = Field(..., alias="workspace-task-id")
+    workspace_task_enforcement_level: TaskEnforcementLevel = Field(
+        ..., alias="workspace-task-enforcement-level"
+    )
+
+    agent_pool_id: Optional[str] = Field(None, alias="agent-pool-id")
+    task_stage: Optional[dict] = Field(None, alias="task-stage")

src/pytfe/models/task_stage.py

@@ -3,14 +3,59 @@
 
 from __future__ import annotations
 
-from pydantic import BaseModel, ConfigDict
+from datetime import datetime
+from enum import Enum
+from typing import List, Optional, TYPE_CHECKING
 
+from pydantic import BaseModel, ConfigDict, Field
 
-# TaskStage represents a HCP Terraform or Terraform Enterprise run's stage where run tasks can occur
+if TYPE_CHECKING:
+    from pytfe.models.task_result import TaskResult
+    from pytfe.models.policy_evaluation import PolicyEvaluation
+
+class Stage(str, Enum):
+    pre_plan = "pre_plan"
+    post_plan = "post_plan"
+    pre_apply = "pre_apply"
+    post_apply = "post_apply"
+
+class TaskStageStatus(str, Enum):
+    pending = "pending"
+    running = "running"
+    passed = "passed"
+    failed = "failed"
+    awaiting_override = "awaiting_override"
+    canceled = "canceled"
+    errored = "errored"
+    unreachable = "unreachable"
+
+class TaskStageStatusTimestamps(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
+
+    errored_at: Optional[datetime] = Field(None, alias="errored-at")
+    running_at: Optional[datetime] = Field(None, alias="running-at")
+    canceled_at: Optional[datetime] = Field(None, alias="canceled-at")
+    failed_at: Optional[datetime] = Field(None, alias="failed-at")
+    passed_at: Optional[datetime] = Field(None, alias="passed-at")
+
+class Permissions(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
+
+    can_override_policy: Optional[bool] = Field(None, alias="can-override-policy")
+    can_override_tasks: Optional[bool] = Field(None, alias="can-override-tasks")
+    can_override: Optional[bool] = Field(None, alias="can-override")
+
+class Actions(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
+
+    is_overridable: Optional[bool] = Field(None, alias="is-overridable")
+
+# TaskStage represents a HCP Terraform or Terraform Enterprise run's stage
 class TaskStage(BaseModel):
     model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
 
     id: str
+
     # stage: Stage = Field(..., alias="stage")
     # status: TaskStageStatus = Field(..., alias="status")
     # status_timestamps: TaskStageStatusTimestamps = Field(..., alias="status-timestamps")
@@ -19,7 +64,28 @@ class TaskStage(BaseModel):
     # permissions: Permissions = Field(..., alias="permissions")
     # actions: Actions = Field(..., alias="actions")
 
+    stage: Stage = Field(..., alias="stage")
+    status: TaskStageStatus = Field(..., alias="status")
+    status_timestamps: TaskStageStatusTimestamps = Field(..., alias="status-timestamps")
+
+    created_at: datetime = Field(..., alias="created-at")
+    updated_at: datetime = Field(..., alias="updated-at")
+
+    permissions: Optional[Permissions] = Field(None, alias="permissions")
+    actions: Optional[Actions] = Field(None, alias="actions")
+
     # # Relations
     # run: Run = Field(..., alias="run")
     # task_results: list[TaskResult] = Field(..., alias="task-results")
     # policy_evaluations: list[PolicyEvaluation] = Field(..., alias="policy-evaluations")
+    
+    run: Optional[dict] = Field(None, alias="run")
+    task_results: Optional[List["TaskResult"]] = Field(None, alias="task-results")
+    policy_evaluations: Optional[List["PolicyEvaluation"]] = Field(
+        None, alias="policy-evaluations"
+    )
+
+from pytfe.models.task_result import TaskResult
+from pytfe.models.policy_evaluation import PolicyEvaluation
+
+TaskStage.model_rebuild()

src/pytfe/resources/task_stage.py

@@ -0,0 +1,70 @@
+# Copyright IBM Corp. 2025, 2026
+# SPDX-License-Identifier: MPL-2.0
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import Any
+
+from ..errors import InvalidRunIDError
+from ..models.task_stage import TaskStage
+from ..utils import _safe_str, valid_string_id
+from ._base import _Service
+
+
+class TaskStages(_Service):
+    """TaskStages provides access to task stage endpoints."""
+
+    # Read
+    def read(self, task_stage_id: str) -> TaskStage:
+        if not valid_string_id(task_stage_id):
+            raise ValueError("Invalid task_stage_id")
+
+        response = self.t.request(
+            "GET",
+            f"/api/v2/task-stages/{task_stage_id}",
+        )
+
+        data = response.json().get("data", {})
+        attributes = data.get("attributes", {})
+        attributes["id"] = _safe_str(data.get("id"))
+
+        return TaskStage.model_validate(attributes)
+
+    # List
+    def list(self, run_id: str) -> Iterator[TaskStage]:
+        if not valid_string_id(run_id):
+            raise InvalidRunIDError()
+
+        path = f"/api/v2/runs/{run_id}/task-stages"
+
+        for item in self._list(path):
+            attributes = item.get("attributes", {})
+            attributes["id"] = item.get("id")
+
+            yield TaskStage.model_validate(attributes)
+
+    # Override
+    def override(
+        self,
+        task_stage_id: str,
+        comment: str | None = None,
+    ) -> TaskStage:
+        if not valid_string_id(task_stage_id):
+            raise ValueError("Invalid task_stage_id")
+
+        body: dict[str, Any] | None = (
+            {"comment": comment} if comment else None
+        )
+
+        response = self.t.request(
+            "POST",
+            f"/api/v2/task-stages/{task_stage_id}/actions/override",
+            json_body=body,
+        )
+
+        data = response.json().get("data", {})
+        attributes = data.get("attributes", {})
+        attributes["id"] = _safe_str(data.get("id"))
+
+        return TaskStage.model_validate(attributes)