[fbsync] Upgrade type hint and others to Python 3.9 (#8814)

Reviewed By: Dan-Flores

Differential Revision: D79175019

fbshipit-source-id: a9b9778af8035fdcb566e02e64730875e4318be9

Co-authored-by: Nicolas Hug <[email protected]>
Co-authored-by: Nicolas Hug <[email protected]>

Author: NicolasHug

test/common_extended_utils.py

@@ -1,7 +1,7 @@
 import os
 from collections import defaultdict
 from numbers import Number
-from typing import Any, List
+from typing import Any
 
 import torch
 from torch.utils._python_dispatch import TorchDispatchMode
@@ -30,7 +30,7 @@ def prod(x):
     return res
 
 
-def matmul_flop(inputs: List[Any], outputs: List[Any]) -> Number:
+def matmul_flop(inputs: list[Any], outputs: list[Any]) -> Number:
     """
     Count flops for matmul.
     """
@@ -43,7 +43,7 @@ def matmul_flop(inputs: List[Any], outputs: List[Any]) -> Number:
     return flop
 
 
-def addmm_flop(inputs: List[Any], outputs: List[Any]) -> Number:
+def addmm_flop(inputs: list[Any], outputs: list[Any]) -> Number:
     """
     Count flops for fully connected layers.
     """
@@ -60,7 +60,7 @@ def addmm_flop(inputs: List[Any], outputs: List[Any]) -> Number:
     return flops
 
 
-def bmm_flop(inputs: List[Any], outputs: List[Any]) -> Number:
+def bmm_flop(inputs: list[Any], outputs: list[Any]) -> Number:
     """
     Count flops for the bmm operation.
     """
@@ -75,9 +75,9 @@ def bmm_flop(inputs: List[Any], outputs: List[Any]) -> Number:
 
 
 def conv_flop_count(
-    x_shape: List[int],
-    w_shape: List[int],
-    out_shape: List[int],
+    x_shape: list[int],
+    w_shape: list[int],
+    out_shape: list[int],
     transposed: bool = False,
 ) -> Number:
     """
@@ -99,7 +99,7 @@ def conv_flop_count(
     return flop
 
 
-def conv_flop(inputs: List[Any], outputs: List[Any]):
+def conv_flop(inputs: list[Any], outputs: list[Any]):
     """
     Count flops for convolution.
     """
@@ -110,7 +110,7 @@ def conv_flop(inputs: List[Any], outputs: List[Any]):
     return conv_flop_count(x_shape, w_shape, out_shape, transposed=transposed)
 
 
-def quant_conv_flop(inputs: List[Any], outputs: List[Any]):
+def quant_conv_flop(inputs: list[Any], outputs: list[Any]):
     """
     Count flops for quantized convolution.
     """
@@ -124,8 +124,8 @@ def transpose_shape(shape):
     return [shape[1], shape[0]] + list(shape[2:])
 
 
-def conv_backward_flop(inputs: List[Any], outputs: List[Any]):
-    grad_out_shape, x_shape, w_shape = [get_shape(i) for i in inputs[:3]]
+def conv_backward_flop(inputs: list[Any], outputs: list[Any]):
+    grad_out_shape, x_shape, w_shape = (get_shape(i) for i in inputs[:3])
     output_mask = inputs[-1]
     fwd_transposed = inputs[7]
     flop_count = 0
@@ -140,7 +140,7 @@ def conv_backward_flop(inputs: List[Any], outputs: List[Any]):
     return flop_count
 
 
-def scaled_dot_product_flash_attention_flop(inputs: List[Any], outputs: List[Any]):
+def scaled_dot_product_flash_attention_flop(inputs: list[Any], outputs: list[Any]):
     # FIXME: this needs to count the flops of this kernel
     # https://github.com/pytorch/pytorch/blob/207b06d099def9d9476176a1842e88636c1f714f/aten/src/ATen/native/cpu/FlashAttentionKernel.cpp#L52-L267
     return 0

test/common_utils.py

@@ -285,7 +285,7 @@ def __init__(
         **other_parameters,
     ):
         if all(isinstance(input, PIL.Image.Image) for input in [actual, expected]):
-            actual, expected = [to_image(input) for input in [actual, expected]]
+            actual, expected = (to_image(input) for input in [actual, expected])
 
         super().__init__(actual, expected, **other_parameters)
         self.mae = mae
@@ -418,7 +418,7 @@ def sample_position(values, max_value):
 
     dtype = dtype or torch.float32
 
-    h, w = [torch.randint(1, s, (num_boxes,)) for s in canvas_size]
+    h, w = (torch.randint(1, s, (num_boxes,)) for s in canvas_size)
     y = sample_position(h, canvas_size[0])
     x = sample_position(w, canvas_size[1])
     r = -360 * torch.rand((num_boxes,)) + 180

test/datasets_utils.py

@@ -15,7 +15,8 @@
 import unittest.mock
 import zipfile
 from collections import defaultdict
-from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union
+from collections.abc import Iterator, Sequence
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
 
@@ -63,6 +64,7 @@ class LazyImporter:
     provide modules listed in MODULES as attributes. They are only imported when accessed.
 
     """
+
     MODULES = (
         "av",
         "lmdb",
@@ -280,7 +282,7 @@ def test_baz(self):
         "download_and_extract_archive",
     }
 
-    def dataset_args(self, tmpdir: str, config: Dict[str, Any]) -> Sequence[Any]:
+    def dataset_args(self, tmpdir: str, config: dict[str, Any]) -> Sequence[Any]:
         """Define positional arguments passed to the dataset.
 
         .. note::
@@ -299,7 +301,7 @@ def dataset_args(self, tmpdir: str, config: Dict[str, Any]) -> Sequence[Any]:
         """
         return (tmpdir,)
 
-    def inject_fake_data(self, tmpdir: str, config: Dict[str, Any]) -> Union[int, Dict[str, Any]]:
+    def inject_fake_data(self, tmpdir: str, config: dict[str, Any]) -> Union[int, dict[str, Any]]:
         """Inject fake data for dataset into a temporary directory.
 
         During the creation of the dataset the download and extract logic is disabled. Thus, the fake data injected
@@ -323,11 +325,11 @@ def inject_fake_data(self, tmpdir: str, config: Dict[str, Any]) -> Union[int, Di
     @contextlib.contextmanager
     def create_dataset(
         self,
-        config: Optional[Dict[str, Any]] = None,
+        config: Optional[dict[str, Any]] = None,
         inject_fake_data: bool = True,
         patch_checks: Optional[bool] = None,
         **kwargs: Any,
-    ) -> Iterator[Tuple[torchvision.datasets.VisionDataset, Dict[str, Any]]]:
+    ) -> Iterator[tuple[torchvision.datasets.VisionDataset, dict[str, Any]]]:
         r"""Create the dataset in a temporary directory.
 
         The configuration passed to the dataset is populated to contain at least all parameters with default values.
@@ -616,11 +618,11 @@ class ImageDatasetTestCase(DatasetTestCase):
     @contextlib.contextmanager
     def create_dataset(
         self,
-        config: Optional[Dict[str, Any]] = None,
+        config: Optional[dict[str, Any]] = None,
         inject_fake_data: bool = True,
         patch_checks: Optional[bool] = None,
         **kwargs: Any,
-    ) -> Iterator[Tuple[torchvision.datasets.VisionDataset, Dict[str, Any]]]:
+    ) -> Iterator[tuple[torchvision.datasets.VisionDataset, dict[str, Any]]]:
         with super().create_dataset(
             config=config,
             inject_fake_data=inject_fake_data,
@@ -799,7 +801,7 @@ def create_image_folder(
     num_examples: int,
     size: Optional[Union[Sequence[int], int, Callable[[int], Union[Sequence[int], int]]]] = None,
     **kwargs: Any,
-) -> List[pathlib.Path]:
+) -> list[pathlib.Path]:
     """Create a folder of random images.
 
     Args:
@@ -821,7 +823,7 @@ def create_image_folder(
     """
     if size is None:
 
-        def size(idx: int) -> Tuple[int, int, int]:
+        def size(idx: int) -> tuple[int, int, int]:
             num_channels = 3
             height, width = torch.randint(3, 11, size=(2,), dtype=torch.int).tolist()
             return (num_channels, height, width)
@@ -913,7 +915,7 @@ def create_video_folder(
     size: Optional[Union[Sequence[int], int, Callable[[int], Union[Sequence[int], int]]]] = None,
     fps=25,
     **kwargs,
-) -> List[pathlib.Path]:
+) -> list[pathlib.Path]:
     """Create a folder of random videos.
 
     Args:

test/test_backbone_utils.py

@@ -1,7 +1,7 @@
 import random
+from collections.abc import Mapping, Sequence
 from copy import deepcopy
 from itertools import chain
-from typing import Mapping, Sequence
 
 import pytest
 import torch

test/test_datasets.py

@@ -14,7 +14,7 @@
 import unittest
 import xml.etree.ElementTree as ET
 import zipfile
-from typing import Callable, Tuple, Union
+from typing import Callable, Union
 
 import datasets_utils
 import numpy as np
@@ -3108,13 +3108,13 @@ class FallingThingsStereoTestCase(datasets_utils.ImageDatasetTestCase):
     FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)))
 
     @staticmethod
-    def _make_dummy_depth_map(root: str, name: str, size: Tuple[int, int]):
+    def _make_dummy_depth_map(root: str, name: str, size: tuple[int, int]):
         file = pathlib.Path(root) / name
         image = np.ones((size[0], size[1]), dtype=np.uint8)
         PIL.Image.fromarray(image).save(file)
 
     @staticmethod
-    def _make_scene_folder(root: str, scene_name: str, size: Tuple[int, int]) -> None:
+    def _make_scene_folder(root: str, scene_name: str, size: tuple[int, int]) -> None:
         root = pathlib.Path(root) / scene_name
         os.makedirs(root, exist_ok=True)
         # jpg images
@@ -3185,7 +3185,7 @@ class SceneFlowStereoTestCase(datasets_utils.ImageDatasetTestCase):
 
     @staticmethod
     def _create_pfm_folder(
-        root: str, name: str, file_name_fn: Callable[..., str], num_examples: int, size: Tuple[int, int]
+        root: str, name: str, file_name_fn: Callable[..., str], num_examples: int, size: tuple[int, int]
     ) -> None:
         root = pathlib.Path(root) / name
         os.makedirs(root, exist_ok=True)
@@ -3268,7 +3268,7 @@ class InStereo2k(datasets_utils.ImageDatasetTestCase):
     ADDITIONAL_CONFIGS = combinations_grid(split=("train", "test"))
 
     @staticmethod
-    def _make_scene_folder(root: str, name: str, size: Tuple[int, int]):
+    def _make_scene_folder(root: str, name: str, size: tuple[int, int]):
         root = pathlib.Path(root) / name
         os.makedirs(root, exist_ok=True)
 

test/test_extended_models.py

@@ -116,7 +116,7 @@ def get_models_from_module(module):
 )
 def test_list_models(module):
     a = set(get_models_from_module(module))
-    b = set(x.replace("quantized_", "") for x in models.list_models(module))
+    b = {x.replace("quantized_", "") for x in models.list_models(module)}
 
     assert len(b) > 0
     assert a == b
@@ -135,7 +135,7 @@ def test_list_models(module):
         ["*resnet*", "*alexnet*"],
         ["*resnet*", "*alexnet*", "*not-existing-model-for-test?"],
         ("*resnet*", "*alexnet*"),
-        set(["*resnet*", "*alexnet*"]),
+        {"*resnet*", "*alexnet*"},
     ],
 )
 @pytest.mark.parametrize(
@@ -151,7 +151,7 @@ def test_list_models(module):
         ["resnet34", "*not-existing-model-for-test?"],
         ["resnet34", "*resnet1*"],
         ("resnet34", "*resnet1*"),
-        set(["resnet34", "*resnet1*"]),
+        {"resnet34", "*resnet1*"},
     ],
 )
 def test_list_models_filters(include_filters, exclude_filters):
@@ -167,15 +167,15 @@ def test_list_models_filters(include_filters, exclude_filters):
         expected = set()
         for include_f in include_filters:
             include_f = include_f.strip("*?")
-            expected = expected | set(x for x in classification_models if include_f in x)
+            expected = expected | {x for x in classification_models if include_f in x}
     else:
         expected = classification_models
 
     if exclude_filters:
         for exclude_f in exclude_filters:
             exclude_f = exclude_f.strip("*?")
             if exclude_f != "":
-                a_exclude = set(x for x in classification_models if exclude_f in x)
+                a_exclude = {x for x in classification_models if exclude_f in x}
                 expected = expected - a_exclude
 
     assert expected == actual
@@ -289,11 +289,11 @@ def test_schema_meta_validation(model_fn):
     bad_names = []
     for w in weights_enum:
         actual_fields = set(w.meta.keys())
-        actual_fields |= set(
+        actual_fields |= {
             ("_metrics", dataset, metric_key)
             for dataset in w.meta.get("_metrics", {}).keys()
             for metric_key in w.meta.get("_metrics", {}).get(dataset, {}).keys()
-        )
+        }
         missing_fields = expected_fields - actual_fields
         unsupported_fields = set(w.meta.keys()) - permitted_fields
         if missing_fields or unsupported_fields:

test/test_functional_tensor.py

@@ -2,8 +2,8 @@
 import itertools
 import math
 import os
+from collections.abc import Sequence
 from functools import partial
-from typing import Sequence
 
 import numpy as np
 import PIL.Image

test/test_image.py

@@ -398,7 +398,7 @@ def test_read_1_bit_png_consistency(shape, mode, tmpdir):
 def test_read_interlaced_png():
     imgs = list(get_images(INTERLACED_PNG, ".png"))
     with Image.open(imgs[0]) as im1, Image.open(imgs[1]) as im2:
-        assert not (im1.info.get("interlace") is im2.info.get("interlace"))
+        assert im1.info.get("interlace") is not im2.info.get("interlace")
     img1 = read_image(imgs[0])
     img2 = read_image(imgs[1])
     assert_equal(img1, img2)
@@ -1040,7 +1040,7 @@ def test_decode_avif_heic_against_pil(decode_fun, mode, pil_mode, filename):
         from torchvision.utils import make_grid
 
         g = make_grid([img, from_pil])
-        F.to_pil_image(g).save((f"/home/nicolashug/out_images/{filename.name}.{pil_mode}.png"))
+        F.to_pil_image(g).save(f"/home/nicolashug/out_images/{filename.name}.{pil_mode}.png")
 
     is_decode_heic = getattr(decode_fun, "__name__", getattr(decode_fun, "name", None)) == "decode_heic"
     if mode == ImageReadMode.RGB and not is_decode_heic:

test/test_onnx.py

@@ -1,6 +1,6 @@
 import io
 from collections import OrderedDict
-from typing import List, Optional, Tuple
+from typing import Optional
 
 import pytest
 import torch
@@ -404,7 +404,7 @@ def forward(self_module, images, features):
             },
         )
 
-    def get_image(self, rel_path: str, size: Tuple[int, int]) -> torch.Tensor:
+    def get_image(self, rel_path: str, size: tuple[int, int]) -> torch.Tensor:
         import os
 
         from PIL import Image
@@ -416,7 +416,7 @@ def get_image(self, rel_path: str, size: Tuple[int, int]) -> torch.Tensor:
 
         return F.convert_image_dtype(F.pil_to_tensor(image))
 
-    def get_test_images(self) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
+    def get_test_images(self) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
         return (
             [self.get_image("encode_jpeg/grace_hopper_517x606.jpg", (100, 320))],
             [self.get_image("fakedata/logos/rgb_pytorch.png", (250, 380))],

test/test_ops.py

@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 from functools import lru_cache
 from itertools import product
-from typing import Callable, List, Tuple
+from typing import Callable
 
 import numpy as np
 import pytest
@@ -100,7 +100,7 @@ def __init__(self, pool: nn.Module):
         super().__init__()
         self.pool = pool
 
-    def forward(self, imgs: Tensor, boxes: List[Tensor]) -> Tensor:
+    def forward(self, imgs: Tensor, boxes: list[Tensor]) -> Tensor:
         return self.pool(imgs, boxes)
 
 
@@ -1485,7 +1485,7 @@ def _run_test(target_fn: Callable, actual_box1, actual_box2, dtypes, atol, expec
             torch.testing.assert_close(out, expected_box, rtol=0.0, check_dtype=False, atol=atol)
 
     @staticmethod
-    def _run_jit_test(target_fn: Callable, actual_box: List):
+    def _run_jit_test(target_fn: Callable, actual_box: list):
         box_tensor = torch.tensor(actual_box, dtype=torch.float)
         expected = target_fn(box_tensor, box_tensor)
         scripted_fn = torch.jit.script(target_fn)