[1/N] Apply UP035 rule in tests (#163947)

Apply UP035 `ruff` rule in tests, but some tests for `fx` and `dynamo` are excluded in case the old typing is the test target.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/163947
Approved by: https://github.com/ezyang

test/distributed/_composable/fsdp/test_fully_shard_comm.py

@@ -6,7 +6,8 @@ import itertools
 import os
 import tempfile
 import unittest
-from typing import Callable, Optional, Union
+from collections.abc import Callable
+from typing import Optional, Union
 from unittest.mock import MagicMock
 
 import torch

test/distributed/_composable/fsdp/test_fully_shard_overlap.py

@@ -3,7 +3,7 @@
 import copy
 import functools
 import unittest
-from typing import Callable
+from collections.abc import Callable
 
 import torch
 import torch.distributed as dist

test/distributed/_composable/test_replicate_with_compiler.py

@@ -3,8 +3,9 @@
 import contextlib
 import functools
 import unittest
+from collections.abc import Callable
 from copy import deepcopy
-from typing import Callable, Optional, Union
+from typing import Optional, Union
 
 import torch
 import torch.distributed as dist

test/distributed/_tools/test_runtime_estimator.py

@@ -1,7 +1,8 @@
 # Owner(s): ["oncall: distributed"]
 import unittest
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Callable, cast, Union
+from typing import Any, cast, Union
 
 import torch
 from torch import nn, optim

test/distributed/checkpoint/test_fsspec.py

@@ -2,8 +2,9 @@
 
 import shutil
 import tempfile
+from collections.abc import Callable
 from functools import wraps
-from typing import Any, Callable, Optional
+from typing import Any, Optional
 
 import torch
 import torch.distributed as dist

test/distributed/checkpoint/test_state_dict.py

@@ -3,8 +3,9 @@
 import copy
 import functools
 import sys
+from collections.abc import Callable
 from itertools import chain
-from typing import Callable, Union
+from typing import Union
 
 import torch
 import torch.distributed as dist

test/distributed/elastic/agent/server/test/local_elastic_agent_test.py

@@ -16,8 +16,9 @@ import tempfile
 import time
 import unittest
 import uuid
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Callable, Optional
+from typing import Optional
 from unittest import mock
 from unittest.mock import Mock, patch
 

test/distributed/elastic/multiprocessing/api_test.py

@@ -15,8 +15,9 @@ import signal
 import sys
 import tempfile
 import time
+from collections.abc import Callable
 from itertools import product
-from typing import Callable, Union
+from typing import Union
 from unittest import mock
 
 import torch

test/distributed/elastic/rendezvous/c10d_rendezvous_backend_test.py

@@ -9,8 +9,9 @@
 import os
 import tempfile
 from base64 import b64encode
+from collections.abc import Callable
 from datetime import timedelta
-from typing import Callable, cast, ClassVar
+from typing import cast, ClassVar
 from unittest import mock, TestCase
 
 from rendezvous_backend_test import RendezvousBackendTestMixin

test/distributed/elastic/rendezvous/dynamic_rendezvous_test.py

@@ -14,8 +14,9 @@ import threading
 import time
 from abc import ABC, abstractmethod
 from base64 import b64encode
+from collections.abc import Callable
 from datetime import datetime, timedelta, timezone
-from typing import Callable, cast, Optional
+from typing import cast, Optional
 from unittest import TestCase
 from unittest.mock import call, MagicMock, Mock, patch, PropertyMock
 

test/distributed/elastic/rendezvous/rendezvous_backend_test.py

@@ -7,7 +7,8 @@
 # LICENSE file in the root directory of this source tree.
 
 from abc import ABC, abstractmethod
-from typing import Any, Callable, cast, Optional
+from collections.abc import Callable
+from typing import Any, cast, Optional
 
 from torch.distributed.elastic.rendezvous import RendezvousStateError
 from torch.distributed.elastic.rendezvous.dynamic_rendezvous import (

test/distributed/fsdp/test_fsdp_core.py

@@ -4,7 +4,8 @@ import functools
 import itertools
 import sys
 import unittest
-from typing import Any, Callable, Optional
+from collections.abc import Callable
+from typing import Any, Optional
 from unittest import mock
 
 import torch

test/distributed/fsdp/test_fsdp_optim_state.py

@@ -2,9 +2,10 @@
 
 import bisect
 import sys
+from collections.abc import Callable
 from copy import deepcopy
 from enum import auto, Enum
-from typing import Any, Callable, Optional
+from typing import Any, Optional
 
 import torch
 import torch.nn as nn

test/distributed/fsdp/test_wrap.py

@@ -5,8 +5,9 @@ import itertools
 import os
 import tempfile
 import unittest
+from collections.abc import Callable
 from enum import auto, Enum
-from typing import Callable, Union
+from typing import Union
 
 import torch
 import torch.nn as nn

test/distributed/pipelining/schedule_registry.py

@@ -2,7 +2,8 @@
 # Owner(s): ["oncall: distributed"]
 # This file is a Schedule zoo for testing torch.distributed.pipelining.
 # It includes schedules designed purely for testing purposes
-from typing import Callable, Optional
+from collections.abc import Callable
+from typing import Optional
 
 from torch.distributed.pipelining.schedules import (
     _Action,

test/distributed/tensor/test_pointwise_ops.py

@@ -1,8 +1,8 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # Owner(s): ["oncall: distributed"]
 
-from collections.abc import Sequence
-from typing import Any, Callable, Optional
+from collections.abc import Callable, Sequence
+from typing import Any, Optional
 from unittest import skip
 
 import torch

test/distributed/tensor/test_xla_integration.py

@@ -3,8 +3,9 @@
 
 import os
 import unittest
+from collections.abc import Callable
 from functools import wraps
-from typing import Any, Callable
+from typing import Any
 
 import numpy as np
 

test/dynamo/test_subclasses.py

@@ -2129,7 +2129,8 @@ class GraphModule(torch.nn.Module):
     @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
     @parametrize("dynamic", [True, False])
     def test_mark_static_with_subclass_desugaring(self, dynamic):
-        from typing import Any, Callable, Optional
+        from collections.abc import Callable
+        from typing import Any, Optional
 
         from torch._dynamo.decorators import mark_static_address
         from torch._inductor.compile_fx import compile_fx

test/functorch/test_aotdispatch.py

@@ -10,9 +10,10 @@ import copy
 import itertools
 import unittest
 import warnings
+from collections.abc import Callable
 from contextlib import ContextDecorator, ExitStack, nullcontext
 from functools import partial, wraps
-from typing import Any, Callable, Optional, Union
+from typing import Any, Optional, Union
 from unittest.mock import patch
 
 from common_utils import (

test/functorch/test_memory_efficient_fusion.py

@@ -3,7 +3,7 @@
 import inspect
 import random
 import unittest
-from typing import Callable
+from collections.abc import Callable
 
 import torch
 import torch.fx as fx

test/inductor/test_aot_inductor_package.py

@@ -9,8 +9,8 @@ import sys
 import tempfile
 import unittest
 import zipfile
+from collections.abc import Callable
 from pathlib import Path
-from typing import Callable
 
 from parameterized import parameterized_class
 

test/inductor/test_cache.py

@@ -8,7 +8,7 @@ from os import environ
 from pathlib import Path
 from random import randint
 from tempfile import gettempdir
-from typing import Any, Callable, Sequence
+from typing import Any, TYPE_CHECKING
 from typing_extensions import Self
 from unittest.mock import patch
 
@@ -20,6 +20,10 @@ from torch.testing._internal.common_utils import (
 )
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+
 class TestMixin:
     @staticmethod
     def abstract_cache_types() -> set[type[icache.Cache]]:

test/inductor/test_cpu_repro.py

@@ -8,7 +8,7 @@ import os
 import platform
 import sys
 import unittest
-from typing import Callable
+from collections.abc import Callable
 from unittest.mock import patch
 
 import torch

test/inductor/test_flex_attention.py

@@ -7,10 +7,11 @@ import string
 import unittest
 import warnings
 from collections import namedtuple
+from collections.abc import Callable
 from contextlib import contextmanager
 from dataclasses import dataclass
 from itertools import product
-from typing import Callable, Optional, TypeVar, Union
+from typing import Optional, TypeVar, Union
 from unittest import expectedFailure, skip, skipUnless
 from unittest.mock import patch
 

test/inductor/test_flex_decoding.py

@@ -5,7 +5,8 @@ import functools
 import sys
 import unittest
 from collections import namedtuple
-from typing import Callable, Optional, Union
+from collections.abc import Callable
+from typing import Optional, Union
 from unittest import expectedFailure
 from unittest.mock import patch
 

test/inductor/test_fxir_backend.py

@@ -6,7 +6,8 @@ Test the FX IR backend.
 import itertools
 import operator
 import unittest
-from typing import Callable, Optional
+from collections.abc import Callable
+from typing import Optional
 
 import sympy
 

test/inductor/test_max_autotune.py

@@ -10,7 +10,8 @@ import random
 import re
 import tempfile
 import unittest
-from typing import Callable, Optional
+from collections.abc import Callable
+from typing import Optional
 from unittest import mock
 
 import torch

test/inductor/test_pattern_matcher.py

@@ -3,7 +3,8 @@ import copy
 import itertools
 import os
 import unittest
-from typing import Callable, Optional
+from collections.abc import Callable
+from typing import Optional
 
 import torch
 import torch._dynamo.config as dynamo_config
@@ -180,8 +181,7 @@ class TestPatternMatcher(TestCase):
             self._test_fused_int_mm_mul_impl(fn2, args, True)
 
     def test_duplicate_search(self):
-        from collections.abc import Iterable
-        from typing import Callable
+        from collections.abc import Callable, Iterable
 
         import torch
         from torch._inductor.pattern_matcher import (

test/inductor/test_profiler.py

@@ -3,7 +3,8 @@ import json
 import os
 import tempfile
 import unittest
-from typing import Callable, Optional
+from collections.abc import Callable
+from typing import Optional
 
 import torch
 import torch._inductor.test_case

test/inductor/test_select_algorithm.py

@@ -2,7 +2,7 @@
 import contextlib
 import functools
 import unittest.mock
-from typing import Callable
+from collections.abc import Callable
 from unittest.mock import patch
 
 import torch

test/inductor/test_torchinductor.py

@@ -19,8 +19,9 @@ import time
 import unittest
 import unittest.mock
 import weakref
+from collections.abc import Callable
 from pathlib import Path
-from typing import Callable, TypeVar
+from typing import TypeVar
 from typing_extensions import ParamSpec
 from unittest.mock import patch
 

test/inductor/test_torchinductor_codegen_config_overrides.py

@@ -1,6 +1,7 @@
 # Owner(s): ["module: inductor"]
 import importlib
-from typing import Any, Callable, Optional
+from collections.abc import Callable
+from typing import Any, Optional
 from unittest import skipIf
 
 import torch

test/inductor/test_torchinductor_strided_blocks.py

@@ -5,7 +5,8 @@ import dataclasses
 import importlib
 import math
 import unittest
-from typing import Any, Callable, Optional, Union
+from collections.abc import Callable
+from typing import Any, Optional, Union
 
 import torch
 import torch.utils._pytree as pytree

test/onnx/onnx_test_common.py

@@ -7,8 +7,8 @@ import dataclasses
 import io
 import os
 import unittest
-from collections.abc import Collection, Iterable, Mapping, Sequence
-from typing import Any, Callable, Optional, Union
+from collections.abc import Callable, Collection, Iterable, Mapping, Sequence
+from typing import Any, Optional, Union
 
 import numpy as np
 import onnxruntime

test/onnx/torchlib/ops_test_common.py

@@ -12,8 +12,8 @@ import pprint
 import sys
 import unittest
 import warnings
-from collections.abc import Collection, Iterable, Mapping, Sequence
-from typing import Any, Callable, Optional, TypeVar
+from collections.abc import Callable, Collection, Iterable, Mapping, Sequence
+from typing import Any, Optional, TypeVar
 
 import error_reproduction
 import numpy as np

test/onnx/torchlib/ops_test_data.py

@@ -39,7 +39,7 @@ from __future__ import annotations
 import copy
 import dataclasses
 import functools
-from typing import Any, Callable, Optional, TYPE_CHECKING
+from typing import Any, Optional, TYPE_CHECKING
 from typing_extensions import Self
 
 import numpy as np
@@ -52,7 +52,7 @@ from torch.testing._internal.opinfo import definitions as opinfo_definitions
 
 
 if TYPE_CHECKING:
-    from collections.abc import Collection
+    from collections.abc import Callable, Collection
 
 
 # Create a copy of the op_db to modify

test/onnx/torchlib/test_ops.py

@@ -25,7 +25,7 @@ errors.
 from __future__ import annotations
 
 import os
-from typing import Callable, Optional, TYPE_CHECKING
+from typing import Optional, TYPE_CHECKING
 
 import error_reproduction
 import numpy as np
@@ -44,7 +44,7 @@ from torch.utils import _pytree as pytree
 
 if TYPE_CHECKING:
     import unittest
-    from collections.abc import Sequence
+    from collections.abc import Callable, Sequence
 
     from torch.testing._internal.opinfo import core as opinfo_core
 

test/profiler/test_memory_profiler.py

@@ -4,8 +4,8 @@ import gc
 import itertools as it
 import textwrap
 import unittest
-from collections.abc import Iterator
-from typing import Callable, Optional
+from collections.abc import Callable, Iterator
+from typing import Optional
 
 import torch
 from torch._C._profiler import _EventType, _TensorMetadata

test/test_public_bindings.py

@@ -7,7 +7,7 @@ import logging
 import os
 import pkgutil
 import unittest
-from typing import Callable
+from collections.abc import Callable
 
 import torch
 from torch._utils_internal import get_file_path_2  # @manual

test/test_sympy_utils.py

@@ -5,7 +5,7 @@ import itertools
 import math
 import pickle
 import sys
-from typing import Callable
+from collections.abc import Callable
 
 import sympy
 

test/test_testing.py

@@ -12,7 +12,8 @@ import re
 import subprocess
 import sys
 import unittest.mock
-from typing import Any, Callable
+from typing import Any
+from collections.abc import Callable
 from collections.abc import Iterator
 
 import torch