astral-sh / ruff

An extremely fast Python linter and code formatter, written in Rust.
https://docs.astral.sh/ruff
MIT License
33.11k stars 1.11k forks source link

PERF401 new preview fixes invalidly hoists extend to list compre #14362

Open Skylion007 opened 1 week ago

Skylion007 commented 1 week ago

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 commented 1 week ago

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 commented 1 week ago

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 commented 1 week ago

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