feat(task-stage): add TaskStage resource and models

examples/task_stage_example.py

@@ -0,0 +1,60 @@
+"""
+Example usage of TaskStages API
+
+Demonstrates:
+- Read a task stage
+- List task stages for a run
+- Override a task stage
+"""
+
+import os
+import sys
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
+
+from pytfe import TFEClient, TFEConfig
+
+
+def main():
+    client = TFEClient(TFEConfig.from_env())
+
+    task_stage_id = os.getenv("TFE_TASK_STAGE_ID")
+    run_id = os.getenv("TFE_RUN_ID")
+
+    if not task_stage_id or not run_id:
+        print("Please set TFE_TASK_STAGE_ID and TFE_RUN_ID")
+        return
+
+    print("=== TaskStages Example ===")
+
+    # READ
+    print("\nReading task stage...")
+    try:
+        stage = client.task_stages.read(task_stage_id)
+        print(f"ID:     {stage.id}")
+        print(f"Stage:  {stage.stage}")
+        print(f"Status: {stage.status}")
+        print(f"Run:    {stage.run.id if stage.run else None}")
+    except Exception as e:
+        print(f"Read failed: {e}")
+
+    # LIST
+    print("\nListing task stages...")
+    try:
+        stages = list(client.task_stages.list(run_id))
+        for s in stages:
+            print(f"{s.id} - {s.status}")
+    except Exception as e:
+        print(f"List failed: {e}")
+
+    # OVERRIDE
+    print("\nOverriding task stage...")
+    try:
+        client.task_stages.override(task_stage_id, comment="Approved")
+        print("Override successful")
+    except Exception as e:
+        print(f"Override failed: {e}")
+
+
+if __name__ == "__main__":
+    main()

src/pytfe/client.py

@@ -45,6 +45,7 @@
 from .resources.state_version_outputs import StateVersionOutputs
 from .resources.state_versions import StateVersions
 from .resources.task_result import TaskResults
+from .resources.task_stage import TaskStages
 from .resources.team import Teams
 from .resources.team_project_access import TeamProjectAccesses
 from .resources.team_token import TeamTokens
@@ -119,6 +120,7 @@ def __init__(self, config: TFEConfig | None = None):
         self.run_task_integrations = RunTaskIntegrations(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.comments = Comments(self._transport)

src/pytfe/errors.py

@@ -490,7 +490,7 @@ def __init__(self, message: str = "must provide at least one policy"):
 
 # Policy Evaluation errors
 class InvalidTaskStageIDError(InvalidValues):
-    """Raised when an invalid task stage ID is provided."""
+    """Raised when a task stage ID is invalid."""
 
     def __init__(self, message: str = "invalid value for task stage ID"):
         super().__init__(message)

src/pytfe/models/__init__.py

@@ -371,6 +371,8 @@
 from .task_result import (
     TaskEnforcementLevel as TaskResultEnforcementLevel,
 )
+
+# ── Task Stage & Task Result ─────────────────────────────────────────────────
 from .task_result import (
     TaskResult,
     TaskResultStatus,
@@ -713,6 +715,9 @@
     "RunEventList",
     "RunEventListOptions",
     "RunEventReadOptions",
+    # Task Stage & Task Result
+    "TaskStage",
+    "TaskResult",
     # Comments
     "Comment",
     "CommentCreateOptions",
@@ -826,14 +831,21 @@
 RegistryProvider.model_rebuild()
 RegistryProviderVersion.model_rebuild()
 RegistryProviderPlatform.model_rebuild()
-
-# Rebuild TaskResult to resolve Run, Workspace, PolicyEvaluation, TaskStage refs
-TaskResult.model_rebuild(
+Run.model_rebuild(
+    raise_errors=False,
+    _types_namespace={"TaskStage": TaskStage},
+)
+TaskStage.model_rebuild(
     raise_errors=False,
     _types_namespace={
-        "PolicyEvaluation": PolicyEvaluation,
         "Run": Run,
+        "TaskResult": TaskResult,
+        "PolicyEvaluation": PolicyEvaluation,
+    },
+)
+TaskResult.model_rebuild(
+    raise_errors=False,
+    _types_namespace={
         "TaskStage": TaskStage,
-        "Workspace": Workspace,
     },
 )

src/pytfe/models/run.py

@@ -5,6 +5,7 @@
 
 from datetime import datetime
 from enum import Enum
+from typing import TYPE_CHECKING
 
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -15,10 +16,12 @@
 from .plan import Plan
 from .policy_check import PolicyCheck
 from .run_event import RunEvent
-from .task_stage import TaskStage
 from .user import User
 from .workspace import Workspace
 
+if TYPE_CHECKING:
+    from .task_stage import TaskStage
+
 
 class RunSource(str, Enum):
     """RunSource represents a source type of a run."""
@@ -327,6 +330,6 @@ class RunDiscardOptions(BaseModel):
 
 
 # Rebuild models to resolve forward references
-Run.model_rebuild()
-RunList.model_rebuild()
-OrganizationRunList.model_rebuild()
+Run.model_rebuild(raise_errors=False)
+RunList.model_rebuild(raise_errors=False)
+OrganizationRunList.model_rebuild(raise_errors=False)

src/pytfe/models/task_result.py

@@ -5,16 +5,12 @@
 
 from datetime import datetime
 from enum import Enum
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from pydantic import BaseModel, ConfigDict, Field
 
 if TYPE_CHECKING:
-    # Imported only for type checking to avoid circular imports.
-    from pytfe.models.policy_evaluation import PolicyEvaluation
-    from pytfe.models.run import Run
     from pytfe.models.task_stage import TaskStage
-    from pytfe.models.workspace import Workspace
 
 
 class TaskResultStatus(str, Enum):
@@ -32,7 +28,7 @@ class TaskEnforcementLevel(str, Enum):
 
 
 class TaskResultStatusTimestamps(BaseModel):
-    model_config = ConfigDict(populate_by_name=True)
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
 
     errored_at: datetime | None = Field(None, alias="errored-at")
     running_at: datetime | None = Field(None, alias="running-at")
@@ -42,16 +38,15 @@ class TaskResultStatusTimestamps(BaseModel):
 
 
 class TaskResult(BaseModel):
-    model_config = ConfigDict(populate_by_name=True)
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
 
     id: str
 
     status: TaskResultStatus | None = Field(None, alias="status")
     message: str | None = Field(None, alias="message")
 
     status_timestamps: TaskResultStatusTimestamps | None = Field(
-        None,
-        alias="status-timestamps",
+        None, alias="status-timestamps"
     )
 
     url: str | None = Field(None, alias="url")
@@ -64,54 +59,10 @@ class TaskResult(BaseModel):
     task_url: str | None = Field(None, alias="task-url")
 
     workspace_task_id: str | None = Field(None, alias="workspace-task-id")
-
     workspace_task_enforcement_level: TaskEnforcementLevel | None = Field(
-        None,
-        alias="workspace-task-enforcement-level",
+        None, alias="workspace-task-enforcement-level"
     )
 
     agent_pool_id: str | None = Field(None, alias="agent-pool-id")
-
-    # Relationships
-    # Forward-referenced to avoid circular imports; resolved lazily below.
+    # relations
     task_stage: TaskStage | None = Field(None, alias="task-stage")
-    run: Run | None = Field(None, alias="run")
-    workspace: Workspace | None = Field(None, alias="workspace")
-    policy_evaluations: list[PolicyEvaluation] | None = Field(
-        None,
-        alias="policy-evaluations",
-    )
-
-    @classmethod
-    def model_validate(cls, *args: Any, **kwargs: Any) -> TaskResult:
-        # Ensure the TaskStage forward reference is resolved before validating.
-        # The import-time rebuild may run while task_stage.py is still
-        # partially loaded (circular import), in which case we retry here.
-        if not getattr(cls, "__pydantic_complete__", True):
-            _rebuild_task_result_model()
-        return super().model_validate(*args, **kwargs)
-
-
-def _rebuild_task_result_model() -> None:
-    # Resolve all forward references once all modules are loaded.
-    try:
-        from pytfe.models.policy_evaluation import PolicyEvaluation
-        from pytfe.models.run import Run
-        from pytfe.models.task_stage import TaskStage
-        from pytfe.models.workspace import Workspace
-
-        TaskResult.model_rebuild(
-            raise_errors=False,
-            _types_namespace={
-                "PolicyEvaluation": PolicyEvaluation,
-                "Run": Run,
-                "TaskStage": TaskStage,
-                "Workspace": Workspace,
-            },
-        )
-    except Exception:
-        # One or more models not yet importable during partial init; safe to skip.
-        pass
-
-
-_rebuild_task_result_model()

src/pytfe/models/task_stage.py

@@ -3,23 +3,82 @@
 
 from __future__ import annotations
 
-from pydantic import BaseModel, ConfigDict
+from datetime import datetime
+from enum import Enum
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from pytfe.models.policy_evaluation import PolicyEvaluation
+from pytfe.models.run import Run
+from pytfe.models.task_result import TaskResult
+
+
+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: datetime | None = Field(None, alias="errored-at")
+    running_at: datetime | None = Field(None, alias="running-at")
+    canceled_at: datetime | None = Field(None, alias="canceled-at")
+    failed_at: datetime | None = Field(None, alias="failed-at")
+    passed_at: datetime | None = Field(None, alias="passed-at")
+
+
+class Permissions(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
+
+    can_override_policy: bool | None = Field(None, alias="can-override-policy")
+    can_override_tasks: bool | None = Field(None, alias="can-override-tasks")
+    can_override: bool | None = Field(None, alias="can-override")
+
+
+class Actions(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
+
+    is_overridable: bool | None = Field(None, alias="is-overridable")
 
 
-# TaskStage represents a HCP Terraform or Terraform Enterprise run's stage where run tasks can occur
 class TaskStage(BaseModel):
     model_config = ConfigDict(populate_by_name=True)
 
     id: str
-    # 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: Permissions = Field(..., alias="permissions")
-    # actions: Actions = Field(..., alias="actions")
-
-    # # Relations
-    # run: Run = Field(..., alias="run")
-    # task_results: list[TaskResult] = Field(..., alias="task-results")
-    # policy_evaluations: list[PolicyEvaluation] = Field(..., alias="policy-evaluations")
+
+    stage: Stage | None = Field(None, alias="stage")
+    status: TaskStageStatus | None = Field(None, alias="status")
+    status_timestamps: TaskStageStatusTimestamps | None = Field(
+        None, alias="status-timestamps"
+    )
+    created_at: datetime | None = Field(None, alias="created-at")
+    updated_at: datetime | None = Field(None, alias="updated-at")
+    permissions: Permissions | None = Field(None, alias="permissions")
+    actions: Actions | None = Field(None, alias="actions")
+
+    # Relationships
+    run: Run | None = Field(None, alias="run")
+    task_results: list[TaskResult] | None = Field(None, alias="task-results")
+    policy_evaluations: list[PolicyEvaluation] | None = Field(
+        None, alias="policy-evaluations"
+    )
+
+
+class TaskStageListOptions(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+
+    page_size: int | None = Field(None, alias="page[size]")

src/pytfe/models/workspace.py

@@ -526,6 +526,7 @@ def _rebuild_workspace_model() -> None:
     """Rebuild Workspace model to resolve forward references."""
     try:
         from .run import Run  # noqa: F401
+        from .task_stage import TaskStage  # noqa: F401
 
         Workspace.model_rebuild()
     except ImportError: