Apply ufmt linter to all py files under tools (pytorch#81285)

With ufmt in place pytorch#81157, we can now use it to gradually format all files. I'm breaking this down into multiple smaller batches to avoid too many merge conflicts later on.

This batch (as copied from the current BLACK linter config):
* `tools/**/*.py`

Upcoming batchs:
* `torchgen/**/*.py`
* `torch/package/**/*.py`
* `torch/onnx/**/*.py`
* `torch/_refs/**/*.py`
* `torch/_prims/**/*.py`
* `torch/_meta_registrations.py`
* `torch/_decomp/**/*.py`
* `test/onnx/**/*.py`

Once they are all formatted, BLACK linter will be removed.

Pull Request resolved: pytorch#81285
Approved by: https://github.com/suo

Author: huydhn

.lintrunner.toml

@@ -661,8 +661,11 @@ init_command = [
 code = 'UFMT'
 # This is a big bang change to reformat every python files in the repo,
 # so it's arguably safer to gradually reformat them in subsequent PR.
-# After that, we can use **/*.py to apply the formatter to everything
-include_patterns = ['tools/linter/**/*.py']
+# After that, we can use **/*.py to apply the formatter to all entries
+# currently enforced by BLACK linter
+include_patterns = [
+    'tools/**/*.py',
+]
 command = [
     'python3',
     'tools/linter/adapters/ufmt_linter.py',

tools/amd_build/build_amd.py

@@ -1,8 +1,8 @@
 #!/usr/bin/env python3
 
 
-import os
 import argparse
+import os
 import sys
 
 sys.path.append(

tools/autograd/context.py

@@ -1,10 +1,10 @@
+import functools
+from typing import Callable
+
 from torchgen.api.autograd import NativeFunctionWithDifferentiabilityInfo as NFWDI
 from torchgen.context import native_function_manager
 from torchgen.utils import T
 
-import functools
-from typing import Callable
-
 # Like tools.api.context.with_native_function, but for
 # NativeFunctionWithDifferentiabilityInfo.
 def with_native_function_with_differentiability_info(

tools/autograd/gen_annotated_fn_args.py

@@ -14,26 +14,28 @@
 torch/testing/_internal/generated
 """
 
-from collections import defaultdict
 import argparse
 import os
 import textwrap
+from collections import defaultdict
 
-from typing import Dict, List, Any
+from typing import Any, Dict, List
 
-from torchgen.gen import parse_native_yaml
-from torchgen.utils import FileManager
+import torchgen.api.python as python
 from torchgen.context import with_native_function
+
+from torchgen.gen import parse_native_yaml
 from torchgen.model import BaseOperatorName, NativeFunction
-import torchgen.api.python as python
+from torchgen.utils import FileManager
+
 from .gen_python_functions import (
-    should_generate_py_binding,
-    is_py_torch_function,
-    is_py_nn_function,
+    is_py_fft_function,
     is_py_linalg_function,
-    is_py_variable_method,
+    is_py_nn_function,
     is_py_special_function,
-    is_py_fft_function,
+    is_py_torch_function,
+    is_py_variable_method,
+    should_generate_py_binding,
 )
 
 

tools/autograd/gen_autograd.py

@@ -25,23 +25,25 @@
 
 import argparse
 import os
+from typing import List
+
 from torchgen.api import cpp
 from torchgen.api.autograd import (
     match_differentiability_info,
     NativeFunctionWithDifferentiabilityInfo,
 )
 from torchgen.gen import parse_native_yaml
 from torchgen.selective_build.selector import SelectiveBuilder
-from typing import List
+
 from . import gen_python_functions
 from .gen_autograd_functions import (
     gen_autograd_functions_lib,
     gen_autograd_functions_python,
 )
-from .gen_trace_type import gen_trace_type
-from .gen_variable_type import gen_variable_type
 from .gen_inplace_or_view_type import gen_inplace_or_view_type
+from .gen_trace_type import gen_trace_type
 from .gen_variable_factories import gen_variable_factories
+from .gen_variable_type import gen_variable_type
 from .load_derivatives import load_derivatives
 
 

tools/autograd/gen_autograd_functions.py

@@ -4,8 +4,6 @@
 #  Functions.h/cpp: subclasses of autograd::Node
 #  python_functions.h/cpp: Python bindings for the above classes
 #
-from .gen_inplace_or_view_type import VIEW_FUNCTIONS
-
 from typing import List, Sequence, Tuple
 
 from torchgen.api.autograd import (
@@ -16,25 +14,27 @@
     uses_single_grad,
 )
 from torchgen.api.types import (
-    Binding,
+    ArrayRefCType,
     BaseCType,
-    OptionalCType,
-    tensorT,
-    longT,
-    doubleT,
-    scalarT,
-    stringT,
+    Binding,
     boolT,
+    doubleT,
     intArrayRefT,
-    tensorListT,
-    MutRefCType,
     ListCType,
-    ArrayRefCType,
+    longT,
+    MutRefCType,
+    OptionalCType,
     optionalIntArrayRefT,
+    scalarT,
+    stringT,
+    tensorListT,
+    tensorT,
 )
 from torchgen.code_template import CodeTemplate
-from torchgen.utils import FileManager
 from torchgen.model import Argument
+from torchgen.utils import FileManager
+
+from .gen_inplace_or_view_type import VIEW_FUNCTIONS
 
 FUNCTION_DECLARATION = CodeTemplate(
     """\

tools/autograd/gen_inplace_or_view_type.py

@@ -4,40 +4,42 @@
 # if updates are needed in torch/csrc/autograd/autograd_not_implemented_fallback.cpp
 # The fallback is expected to mimick this codegen, so we should keep the two in sync.
 
+from typing import Dict, List, Optional, Sequence, Tuple
+
 from torchgen.api import cpp
 from torchgen.api.autograd import (
-    NativeFunctionWithDifferentiabilityInfo,
-    gen_differentiable_outputs,
     dispatch_strategy,
+    gen_differentiable_outputs,
+    NativeFunctionWithDifferentiabilityInfo,
 )
 from torchgen.api.types import (
-    Binding,
-    DispatcherSignature,
-    CType,
     BaseCType,
-    OptionalCType,
-    longT,
+    Binding,
     boolT,
+    CType,
+    DispatcherSignature,
     intArrayRefT,
+    longT,
+    OptionalCType,
     symIntArrayRefT,
 )
 from torchgen.code_template import CodeTemplate
 from torchgen.context import with_native_function
 from torchgen.model import (
-    Type,
     NativeFunction,
+    SchemaKind,
     SelfArgument,
     TensorOptionsArguments,
-    SchemaKind,
+    Type,
 )
-from typing import List, Optional, Sequence, Tuple, Dict
 from torchgen.utils import FileManager
+
 from .context import with_native_function_with_differentiability_info
 from .gen_trace_type import (
+    get_return_value,
     MANUAL_AUTOGRAD,
-    type_wrapper_name,
     tie_return_values,
-    get_return_value,
+    type_wrapper_name,
 )
 
 # See NOTE [ Autograd View Variables ] in variable.h for details.

tools/autograd/gen_python_functions.py

@@ -30,22 +30,15 @@
 #   message, but use what's there
 #
 
-from collections import defaultdict
 import itertools
 import re
-import yaml
+from collections import defaultdict
 
-from .gen_trace_type import should_trace
+from typing import Callable, Dict, List, Optional, Sequence, Set, Tuple
 
-from torchgen.code_template import CodeTemplate
+import yaml
 from torchgen.api import cpp
-from torchgen.api.types import CppSignatureGroup
 from torchgen.api.python import (
-    PythonArgument,
-    PythonSignature,
-    PythonSignatureDeprecated,
-    PythonSignatureGroup,
-    PythonSignatureNativeFunctionPair,
     arg_parser_output_exprs,
     argument_type_str,
     cpp_dispatch_exprs,
@@ -55,20 +48,22 @@
     dispatch_lambda_return_str,
     has_tensor_options,
     namedtuple_fieldnames,
+    PythonArgument,
+    PythonSignature,
+    PythonSignatureDeprecated,
+    PythonSignatureGroup,
+    PythonSignatureNativeFunctionPair,
     signature,
 )
-from torchgen.gen import cpp_string, parse_native_yaml, parse_tags_yaml
+from torchgen.api.types import CppSignatureGroup
+
+from torchgen.code_template import CodeTemplate
 from torchgen.context import with_native_function
-from torchgen.model import (
-    Argument,
-    BaseOperatorName,
-    NativeFunction,
-    Type,
-    Variant,
-)
-from torchgen.utils import split_name_params, YamlLoader, FileManager
+from torchgen.gen import cpp_string, parse_native_yaml, parse_tags_yaml
+from torchgen.model import Argument, BaseOperatorName, NativeFunction, Type, Variant
+from torchgen.utils import FileManager, split_name_params, YamlLoader
 
-from typing import Dict, Optional, List, Tuple, Set, Sequence, Callable
+from .gen_trace_type import should_trace
 
 #
 # declarations blocklist

tools/autograd/gen_trace_type.py

@@ -1,17 +1,13 @@
 import itertools
-from typing import List, Sequence, Union, Dict
+from typing import Dict, List, Sequence, Union
 
-from torchgen.api.types import DispatcherSignature
 from torchgen.api import cpp
+
+from torchgen.api.types import DispatcherSignature
 from torchgen.code_template import CodeTemplate
 from torchgen.context import with_native_function
+from torchgen.model import Argument, NativeFunction, SchemaKind, TensorOptionsArguments
 from torchgen.utils import FileManager
-from torchgen.model import (
-    Argument,
-    NativeFunction,
-    SchemaKind,
-    TensorOptionsArguments,
-)
 
 # Note [Manual Backend kernels]
 # For these ops, we want to manually register to dispatch key Backend and

tools/autograd/gen_variable_factories.py

@@ -3,15 +3,16 @@
 # This writes one file: variable_factories.h
 
 import re
-from typing import Optional, List
+from typing import List, Optional
 
-from torchgen.api.types import CppSignatureGroup
-from torchgen.api import cpp
 import torchgen.api.python as python
-from torchgen.gen import parse_native_yaml
+from torchgen.api import cpp
+
+from torchgen.api.types import CppSignatureGroup
 from torchgen.context import with_native_function
-from torchgen.utils import mapMaybe, FileManager
+from torchgen.gen import parse_native_yaml
 from torchgen.model import NativeFunction, TensorOptionsArguments, Variant
+from torchgen.utils import FileManager, mapMaybe
 
 OPTIONAL_TYPE_PATTERN = re.compile(r"c10::optional<(.+)>")
 TYPE_PATTERN = re.compile(r"(?:const\s+)?([A-Z]\w+)")