PERF401 new preview fixes invalidly hoists extend to list compre

[Skylion007]

I tried running the new autofixes on the PyTorch codebase and was mostly impressed, but found one really annoying (and handable edge case).

Here is an example bad fix

     @dist_init
     def test_wait_all_with_exception(self):
-        futs = []
+        futs = [rpc.rpc_async(dst, raise_func) for _ in range(10)]
         dst = worker_name((self.rank + 1) % self.world_size)
-        for _ in range(10):
-            futs.append(rpc.rpc_async(dst, raise_func))

         with self.assertRaisesRegex(ValueError, "Expected error"):
             torch.futures.wait_all(futs)

Here is an example diff generated by ruff. Note that list comprehensions uses dst even though dst is first defined on the line underneath the list comprehension. Ruff can already detect this because it immeaditely created a bunch of ruff F821 errors as soon as the fixes were applied. It would be good not to hoist the forloop from an extend to a list comprehensions if there are any variables needed for the list comprehension defined or mutated in anyway. I was kind of surprised given that it did properly not hoist the function if there were any comments in between the list definition and the loop.

ruff 0.7.4
ruff check --select=PERF401 --fix --unsafe-fixes --preview

FYI @w0nder1ng

[w0nder1ng]

I didn't consider that when I was writing the fix, and I guess the question now is how this case should be handled. I guess a possible fix could look like:

     @dist_init
     def test_wait_all_with_exception(self):
-        futs = []
         dst = worker_name((self.rank + 1) % self.world_size)
-        for _ in range(10):
-            futs.append(rpc.rpc_async(dst, raise_func))
+        futs = [rpc.rpc_async(dst, raise_func) for _ in range(10)]
         with self.assertRaisesRegex(ValueError, "Expected error"):
             torch.futures.wait_all(futs)

As long as we check that the futs variable isn't used between futs = [] and the for loop, it should be fine.

I was kind of surprised given that it did properly not hoist the function if there were any comments in between the list definition and the loop.

If you have an example of this, I'd be happy to take a look.

[Skylion007]

Sorry typo, it properly did NOT hoist the function if there were comments (to preserve the comments). As it did not hoist the list comprehension if there was a comment there. It did do if there was code in the way, which was surprising.

[Skylion007]

Also, FYI once these fixes are landed. We should look into PERF403, as the fixes will share a lot of similar logic. :)