ci: fix type errors resulting from pyright upgrade

Author: ahal

src/taskgraph/generator.py

@@ -6,14 +6,14 @@
 import logging
 import os
 from dataclasses import dataclass
-from typing import Dict
+from typing import Callable, Dict, Optional, Union
 
 from . import filter_tasks
 from .config import GraphConfig, load_graph_config
 from .graph import Graph
 from .morph import morph
 from .optimize.base import optimize_task_graph
-from .parameters import parameters_loader
+from .parameters import Parameters, parameters_loader
 from .task import Task
 from .taskgraph import TaskGraph
 from .transforms.base import TransformConfig, TransformSequence
@@ -119,10 +119,10 @@ class TaskGraphGenerator:
 
     def __init__(
         self,
-        root_dir,
-        parameters,
-        decision_task_id="DECISION-TASK",
-        write_artifacts=False,
+        root_dir: Optional[str],
+        parameters: Union[Parameters, Callable[[GraphConfig], Parameters]],
+        decision_task_id: str = "DECISION-TASK",
+        write_artifacts: bool = False,
     ):
         """
         @param root_dir: root directory containing the Taskgraph config.yml file

src/taskgraph/optimize/base.py

@@ -15,6 +15,7 @@
 import logging
 from abc import ABCMeta, abstractmethod
 from collections import defaultdict
+from typing import Dict, Set
 
 from slugid import nice as slugid
 
@@ -365,11 +366,11 @@ def replace_tasks(
 
 
 def get_subgraph(
-    target_task_graph,
-    removed_tasks,
-    replaced_tasks,
-    label_to_taskid,
-    decision_task_id,
+    target_task_graph: TaskGraph,
+    removed_tasks: Set[str],
+    replaced_tasks: Set[str],
+    label_to_taskid: Dict[str, str],
+    decision_task_id: str,
 ):
     """
     Return the subgraph of target_task_graph consisting only of
@@ -424,14 +425,15 @@ def get_subgraph(
                 }
             )
 
+        assert task.task_id
         task.task = resolve_task_references(
             task.label,
             task.task,
             task_id=task.task_id,
             decision_task_id=decision_task_id,
             dependencies=named_task_dependencies,
         )
-        deps = task.task.setdefault("dependencies", [])  # type: ignore
+        deps = task.task.setdefault("dependencies", [])
         deps.extend(sorted(named_task_dependencies.values()))
         tasks_by_taskid[task.task_id] = task
 

src/taskgraph/transforms/run/common.py

@@ -170,11 +170,8 @@ def support_caches(
     if use_caches is None:
         # Use project default values for filtering caches, default to
         # checkout cache if no selection is specified.
-        use_caches = (
-            config.graph_config.get("taskgraph", {})
-            .get("run", {})
-            .get("use-caches", ["checkout"])
-        )
+        taskgraph_config = config.graph_config.get("taskgraph") or {}
+        use_caches = taskgraph_config.get("run", {}).get("use-caches", ["checkout"])
 
     for name, cache_cfg in CACHES.items():
         if not should_use_cache(name, use_caches, run["checkout"]):

src/taskgraph/util/parameterization.py

@@ -4,6 +4,7 @@
 
 
 import re
+from typing import Any, Dict
 
 from taskgraph.util.taskcluster import get_artifact_url
 from taskgraph.util.time import json_time_from_now
@@ -44,7 +45,13 @@ def resolve_timestamps(now, task_def):
     )
 
 
-def resolve_task_references(label, task_def, task_id, decision_task_id, dependencies):
+def resolve_task_references(
+    label: str,
+    task_def: Dict[str, Any],
+    task_id: str,
+    decision_task_id: str,
+    dependencies: Dict[str, str],
+) -> Dict[str, Any]:
     """Resolve all instances of ``{'task-reference': '..<..>..'} ``
     and ``{'artifact-reference`: '..<dependency/artifact/path>..'}``
     in the given task definition, using the given dependencies.
@@ -91,10 +98,12 @@ def repl(match):
 
         return ARTIFACT_REFERENCE_PATTERN.sub(repl, val)
 
-    return _recurse(
+    result = _recurse(
         task_def,
         {
             "task-reference": task_reference,
             "artifact-reference": artifact_reference,
         },
     )
+    assert isinstance(result, dict)
+    return result

test/test_util_parameterization.py

@@ -211,8 +211,8 @@ def test_artifact_refs_private(monkeypatch, assert_artifact_refs):
     monkeypatch.setenv("TASKCLUSTER_PROXY_URL", tc_proxy_url)
 
     assert_artifact_refs(
-        {"artifact-reference": "<edge1/private/foo>"},
-        f"{tc_proxy_url}/api/queue/v1/task/tid1/artifacts/private/foo",
+        {"url": {"artifact-reference": "<edge1/private/foo>"}},
+        {"url": f"{tc_proxy_url}/api/queue/v1/task/tid1/artifacts/private/foo"},
     )
 
 
@@ -255,6 +255,10 @@ def test_artifact_refs_badly_formed():
     "resolve_task_references ignores badly-formatted artifact references"
     for inv in ["<edge1>", "edge1/foo>", "<edge1>/foo", "<edge1>foo"]:
         resolved = resolve_task_references(
-            "subject", {"artifact-reference": inv}, "tid-self", "tid-decision", {}
+            "subject",
+            {"key": {"artifact-reference": inv}},
+            "tid-self",
+            "tid-decision",
+            {},
         )
-        assert resolved == inv
+        assert resolved == {"key": inv}