[BE][Easy] enable postponed annotations in tools (pytorch#129375)

Pull Request resolved: pytorch#129375
Approved by: https://github.com/malfet

Author: XuehaiPan

tools/alerts/create_alerts.py

@@ -1,16 +1,19 @@
 #!/usr/bin/env python3
 
+from __future__ import annotations
+
 import argparse
 import json
 import os
 import re
 from collections import defaultdict
 from difflib import SequenceMatcher
-from typing import Any, Dict, List, Set, Tuple
+from typing import Any
 
 import requests
 from setuptools import distutils  # type: ignore[import]
 
+
 ALL_SKIPPED_THRESHOLD = 100
 SIMILARITY_THRESHOLD = 0.75
 FAILURE_CHAIN_THRESHOLD = 2
@@ -65,14 +68,14 @@
 
 class JobStatus:
     job_name: str = ""
-    jobs: List[Any] = []
+    jobs: list[Any] = []
     current_status: Any = None
-    job_statuses: List[Any] = []
-    filtered_statuses: List[Any] = []
-    failure_chain: List[Any] = []
-    flaky_jobs: List[Any] = []
+    job_statuses: list[Any] = []
+    filtered_statuses: list[Any] = []
+    failure_chain: list[Any] = []
+    flaky_jobs: list[Any] = []
 
-    def __init__(self, job_name: str, job_statuses: List[Any]):
+    def __init__(self, job_name: str, job_statuses: list[Any]) -> None:
         self.job_name = job_name
         self.job_statuses = job_statuses
 
@@ -93,7 +96,7 @@ def get_current_status(self) -> Any:
                 return status
         return None
 
-    def get_unique_failures(self, jobs: List[Any]) -> Dict[str, List[Any]]:
+    def get_unique_failures(self, jobs: list[Any]) -> dict[str, list[Any]]:
         """
         Returns list of jobs grouped by failureCaptures from the input list
         """
@@ -120,7 +123,7 @@ def get_unique_failures(self, jobs: List[Any]) -> Dict[str, List[Any]]:
         return failures
 
     # A flaky job is if it's the only job that has that failureCapture and is not the most recent job
-    def get_flaky_jobs(self) -> List[Any]:
+    def get_flaky_jobs(self) -> list[Any]:
         unique_failures = self.get_unique_failures(self.filtered_statuses)
         flaky_jobs = []
         for failure in unique_failures:
@@ -134,7 +137,7 @@ def get_flaky_jobs(self) -> List[Any]:
 
     # The most recent failure chain is an array of jobs that have the same-ish failures.
     # A success in the middle of the chain will terminate the chain.
-    def get_most_recent_failure_chain(self) -> List[Any]:
+    def get_most_recent_failure_chain(self) -> list[Any]:
         failures = []
         found_most_recent_failure = False
 
@@ -178,7 +181,7 @@ def fetch_hud_data(repo: str, branch: str) -> Any:
 
 
 # Creates a Dict of Job Name -> [JobData]. Essentially a Column in HUD
-def map_job_data(jobNames: Any, shaGrid: Any) -> Dict[str, Any]:
+def map_job_data(jobNames: Any, shaGrid: Any) -> dict[str, Any]:
     jobData = defaultdict(list)
     for sha in shaGrid:
         for ind, job in enumerate(sha["jobs"]):
@@ -196,13 +199,13 @@ def is_job_skipped(job: Any) -> bool:
     return conclusion in (NEUTRAL, SKIPPED) or conclusion is None
 
 
-def get_failed_jobs(job_data: List[Any]) -> List[Any]:
+def get_failed_jobs(job_data: list[Any]) -> list[Any]:
     return [job for job in job_data if job["conclusion"] == "failure"]
 
 
 def classify_jobs(
-    all_job_names: List[str], sha_grid: Any, filtered_jobs_names: Set[str]
-) -> Tuple[List[JobStatus], List[Any]]:
+    all_job_names: list[str], sha_grid: Any, filtered_jobs_names: set[str]
+) -> tuple[list[JobStatus], list[Any]]:
     """
     Creates Job Statuses which has the logic for if need to alert or if there's flaky jobs.
     Classifies jobs into jobs to alert on and flaky jobs.
@@ -212,7 +215,7 @@ def classify_jobs(
     :return:
     """
     job_data = map_job_data(all_job_names, sha_grid)
-    job_statuses: List[JobStatus] = []
+    job_statuses: list[JobStatus] = []
     for job in job_data:
         job_statuses.append(JobStatus(job, job_data[job]))
 
@@ -230,7 +233,7 @@ def classify_jobs(
 
 
 # filter job names that don't match the regex
-def filter_job_names(job_names: List[str], job_name_regex: str) -> List[str]:
+def filter_job_names(job_names: list[str], job_name_regex: str) -> list[str]:
     if job_name_regex:
         return [
             job_name for job_name in job_names if re.match(job_name_regex, job_name)
@@ -240,7 +243,7 @@ def filter_job_names(job_names: List[str], job_name_regex: str) -> List[str]:
 
 def get_recurrently_failing_jobs_alerts(
     repo: str, branch: str, job_name_regex: str
-) -> List[Dict[str, Any]]:
+) -> list[dict[str, Any]]:
     job_names, sha_grid = fetch_hud_data(repo=repo, branch=branch)
 
     filtered_job_names = set(filter_job_names(job_names, job_name_regex))

tools/autograd/gen_annotated_fn_args.py

@@ -14,18 +14,17 @@
 torch/testing/_internal/generated
 """
 
+from __future__ import annotations
+
 import argparse
 import os
 import textwrap
 from collections import defaultdict
-
-from typing import Any, Dict, List, Sequence
+from typing import Any, Sequence, TYPE_CHECKING
 
 import torchgen.api.python as python
 from torchgen.context import with_native_function
-
 from torchgen.gen import parse_native_yaml
-from torchgen.model import Argument, BaseOperatorName, NativeFunction
 from torchgen.utils import FileManager
 
 from .gen_python_functions import (
@@ -39,6 +38,10 @@
 )
 
 
+if TYPE_CHECKING:
+    from torchgen.model import Argument, BaseOperatorName, NativeFunction
+
+
 def gen_annotated(
     native_yaml_path: str, tags_yaml_path: str, out: str, autograd_dir: str
 ) -> None:
@@ -53,9 +56,9 @@ def gen_annotated(
         (is_py_fft_function, "torch._C._fft"),
         (is_py_variable_method, "torch.Tensor"),
     )
-    annotated_args: List[str] = []
+    annotated_args: list[str] = []
     for pred, namespace in mappings:
-        groups: Dict[BaseOperatorName, List[NativeFunction]] = defaultdict(list)
+        groups: dict[BaseOperatorName, list[NativeFunction]] = defaultdict(list)
         for f in native_functions:
             if not should_generate_py_binding(f) or not pred(f):
                 continue
@@ -77,7 +80,7 @@ def gen_annotated(
 
 @with_native_function
 def gen_annotated_args(f: NativeFunction) -> str:
-    def _get_kwargs_func_exclusion_list() -> List[str]:
+    def _get_kwargs_func_exclusion_list() -> list[str]:
         # functions that currently don't work with kwargs in test_overrides.py
         return [
             "diagonal",
@@ -87,12 +90,12 @@ def _get_kwargs_func_exclusion_list() -> List[str]:
         ]
 
     def _add_out_arg(
-        out_args: List[Dict[str, Any]], args: Sequence[Argument], *, is_kwarg_only: bool
+        out_args: list[dict[str, Any]], args: Sequence[Argument], *, is_kwarg_only: bool
     ) -> None:
         for arg in args:
             if arg.default is not None:
                 continue
-            out_arg: Dict[str, Any] = {}
+            out_arg: dict[str, Any] = {}
             out_arg["is_kwarg_only"] = str(is_kwarg_only)
             out_arg["name"] = arg.name
             out_arg["simple_type"] = python.argument_type_str(
@@ -103,7 +106,7 @@ def _add_out_arg(
                 out_arg["size"] = size_t
             out_args.append(out_arg)
 
-    out_args: List[Dict[str, Any]] = []
+    out_args: list[dict[str, Any]] = []
     _add_out_arg(out_args, f.func.arguments.flat_positional, is_kwarg_only=False)
     if f"{f.func.name.name}" not in _get_kwargs_func_exclusion_list():
         _add_out_arg(out_args, f.func.arguments.flat_kwarg_only, is_kwarg_only=True)

tools/autograd/gen_autograd.py

@@ -22,9 +22,10 @@
 #  gen_python_functions.py: generates Python bindings to THPVariable
 #
 
+from __future__ import annotations
+
 import argparse
 import os
-from typing import List
 
 from torchgen.api import cpp
 from torchgen.api.autograd import (
@@ -69,7 +70,7 @@ def gen_autograd(
         ),
         key=lambda f: cpp.name(f.func),
     )
-    fns_with_diff_infos: List[
+    fns_with_diff_infos: list[
         NativeFunctionWithDifferentiabilityInfo
     ] = match_differentiability_info(fns, differentiability_infos)
 

tools/autograd/gen_autograd_functions.py

@@ -4,7 +4,10 @@
 #  Functions.h/cpp: subclasses of autograd::Node
 #  python_functions.h/cpp: Python bindings for the above classes
 #
-from typing import Dict, List, Sequence, Tuple
+
+from __future__ import annotations
+
+from typing import Sequence
 
 from torchgen.api.autograd import (
     Derivative,
@@ -43,6 +46,7 @@
 
 from .gen_inplace_or_view_type import VIEW_FUNCTIONS
 
+
 FUNCTION_DECLARATION = CodeTemplate(
     """\
 #ifdef _WIN32
@@ -443,8 +447,8 @@
 
 
 def get_infos_with_derivatives_list(
-    differentiability_infos: Dict[FunctionSchema, Dict[str, DifferentiabilityInfo]]
-) -> List[DifferentiabilityInfo]:
+    differentiability_infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]]
+) -> list[DifferentiabilityInfo]:
     diff_info_list = [
         info
         for diffinfo_dict in differentiability_infos.values()
@@ -456,7 +460,7 @@ def get_infos_with_derivatives_list(
 
 def gen_autograd_functions_lib(
     out: str,
-    differentiability_infos: Dict[FunctionSchema, Dict[str, DifferentiabilityInfo]],
+    differentiability_infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]],
     template_path: str,
 ) -> None:
     """Functions.h and Functions.cpp body
@@ -490,7 +494,7 @@ def gen_autograd_functions_lib(
 
 def gen_autograd_functions_python(
     out: str,
-    differentiability_infos: Dict[FunctionSchema, Dict[str, DifferentiabilityInfo]],
+    differentiability_infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]],
     template_path: str,
 ) -> None:
     fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)
@@ -536,17 +540,17 @@ def gen_autograd_functions_python(
 
 
 def process_function(info: DifferentiabilityInfo, template: CodeTemplate) -> str:
-    saved_variables: List[str] = []
-    release_variables: List[str] = []
-    saved_list_sizes: List[str] = []
-    unpack: List[str] = []
-    asserts: List[str] = []
-    compute_index_ranges: List[str] = []
-    getter_definitions: List[str] = []
-    py_getsetdef_structs: List[str] = []
-    compiled_args: List[str] = []
-    apply_with_saved_before: List[str] = []
-    apply_with_saved_after: List[str] = []
+    saved_variables: list[str] = []
+    release_variables: list[str] = []
+    saved_list_sizes: list[str] = []
+    unpack: list[str] = []
+    asserts: list[str] = []
+    compute_index_ranges: list[str] = []
+    getter_definitions: list[str] = []
+    py_getsetdef_structs: list[str] = []
+    compiled_args: list[str] = []
+    apply_with_saved_before: list[str] = []
+    apply_with_saved_after: list[str] = []
 
     for arg in info.args_with_derivatives:
         if arg.type in TENSOR_LIST_LIKE_CTYPES:
@@ -807,7 +811,7 @@ def save_var(var: SavedAttribute, is_output: bool) -> None:
     else:
         will_release_variables = ""
 
-    body: List[str] = []
+    body: list[str] = []
 
     if uses_single_grad(info):
         body.append("const auto& grad = grads[0];")
@@ -821,7 +825,7 @@ def save_var(var: SavedAttribute, is_output: bool) -> None:
     def emit_derivative(
         derivative: Derivative,
         args_with_derivatives: Sequence[Binding],
-    ) -> Tuple[bool, str]:
+    ) -> tuple[bool, str]:
         formula = derivative.formula
         var_names = derivative.var_names
         if len(var_names) == 1:
@@ -857,7 +861,7 @@ def emit_derivative(
             else:
                 grad_input_mask = ""
             idx_ranges = ", ".join(f"{n}_ix" for n in var_names)
-            copy_ranges: List[str] = []
+            copy_ranges: list[str] = []
             for i, n in enumerate(var_names):
                 copy_ranges.append(DERIVATIVE_MULTI_COPY_RANGE.substitute(name=n, i=i))
             return False, DERIVATIVE_MULTI.substitute(