Repository: facebookexperimental/usort
Pull request #161 (2022-05-17 23:46:30)

Optimize passing multiple paths to usort CLI
Previously, when passing multiple paths to the usort CLI, it would
call `usort_path()` separately for each path given, preventing any
chance of parallel processing. This means `usort check usort/**/*.py`
would take ~10X longer than `usort check usort`, and also results in
double formatting if a path appeared or was a child of multiple
arguments:

```
(.venv) jreese@butterfree ~/workspace/usort main  » time usort check usort
usort check usort  3.64s user 0.32s system 338% cpu 1.169 total

(.venv) jreese@butterfree ~/workspace/usort main  » time usort check usort usort usort
usort check usort usort usort  10.67s user 0.89s system 470% cpu 2.458 total

(.venv) jreese@butterfree ~/workspace/usort main  » time usort check usort/**/*.py usort/**/*.py usort/**/*.py
usort check usort/**/*.py usort/**/*.py usort/**/*.py  13.41s user 1.59s system 62% cpu 24.141 total
```

With the changes from #160, this is partially masked by removing the
cost of starting child processes, but still results in linear time
increase when passing many arguments:

```
(.venv) jreese@butterfree ~/workspace/usort patch-1  » time usort check usort
usort check usort  3.68s user 0.36s system 318% cpu 1.267 total

(.venv) jreese@butterfree ~/workspace/usort patch-1  » time usort check usort usort usort
usort check usort usort usort  10.37s user 0.86s system 483% cpu 2.323 total

(.venv) jreese@butterfree ~/workspace/usort patch-1  » time usort check usort/**/*.py usort/**/*.py usort/**/*.py
usort check usort/**/*.py usort/**/*.py usort/**/*.py  2.45s user 0.06s system 89% cpu 2.803 total
```

This builds on #160 by updating `usort_path()` to optionally take an
iterable of Path objects, and materializes a set of paths to sort before
actually starting child processes. This allows maximum utilization of
multiprocessing, and also eliminates any double-formatting if the CLI
arguments include any overlap of paths:

```
(.venv) jreese@butterfree ~/workspace/usort gather  ‹1› » time usort check usort
usort check usort  3.82s user 0.36s system 333% cpu 1.255 total

(.venv) jreese@butterfree ~/workspace/usort gather  » time usort check usort usort usort
usort check usort usort usort  3.88s user 0.37s system 356% cpu 1.192 total

(.venv) jreese@butterfree ~/workspace/usort gather  » time usort check usort/**/*.py usort/**/*.py usort/**/*.py
usort check usort/**/*.py usort/**/*.py usort/**/*.py  3.68s user 0.36s system 323% cpu 1.249 total
```


Patch:
diff --git a/usort/api.py b/usort/api.py
index 32d14bc..0b23edd 100644
--- a/usort/api.py
+++ b/usort/api.py
@@ -6,7 +6,7 @@
 import sys
 from functools import partial
 from pathlib import Path
-from typing import Iterable, Optional, Tuple
+from typing import Iterable, Optional, Set, Tuple, Union
 from warnings import warn
 
 from trailrunner import run, walk
@@ -117,7 +117,9 @@ def usort_file(path: Path, *, write: bool = False) -> Result:
         )
 
 
-def usort_path(path: Path, *, write: bool = False) -> Iterable[Result]:
+def usort_path(
+    paths: Union[Path, Iterable[Path]], write: bool = False
+) -> Iterable[Result]:
     """
     For a given path, format it, or any python files in it, and yield :class:`Result` s.
 
@@ -125,16 +127,25 @@ def usort_path(path: Path, *, write: bool = False) -> Iterable[Result]:
     excluding any files or directories that match the project root's ``.gitignore`` or
     any configured :py:attr:`excludes` patterns in the associated ``pyproject.toml``.
     """
-    with timed(f"total for {path}"):
-        with timed(f"walking {path}"):
-            config = Config.find(path)
-            paths = list(walk(path, excludes=config.excludes))
+    if isinstance(paths, Path):
+        source_paths: Iterable[Path] = [paths]
+    else:
+        source_paths = paths
+
+    with timed("total"):
+        materialized_paths: Set[Path] = set()
+
+        for path in source_paths:
+            with timed(f"walking {path}"):
+                config = Config.find(path)
+                materialized_paths.update(walk(path, excludes=config.excludes))
 
         fn = partial(usort_file, write=write)
-        if len(paths) == 1:
-            results = [fn(paths[0])]  # shave off multiprocessing overhead
+        if len(materialized_paths) == 1:
+            # shave off multiprocessing overhead
+            results = [fn(materialized_paths.pop())]
         else:
-            results = [v for v in run(paths, fn).values()]
+            results = [v for v in run(materialized_paths, fn).values()]
         return results
 
 
diff --git a/usort/cli.py b/usort/cli.py
index 2ddf58c..edca116 100644
--- a/usort/cli.py
+++ b/usort/cli.py
@@ -117,21 +117,20 @@ def check(filenames: List[str]) -> int:
         raise click.ClickException("Provide some filenames")
 
     return_code = 0
-    for f in filenames:
-        path = Path(f)
-        for result in usort_path(path, write=False):
-            if result.error:
-                click.echo(f"Error sorting {result.path}: {result.error}")
-                return_code |= 1
+    paths = [Path(f) for f in filenames]
+    for result in usort_path(paths, write=False):
+        if result.error:
+            click.echo(f"Error sorting {result.path}: {result.error}")
+            return_code |= 1
 
-            for warning in result.warnings:
-                click.echo(f"Warning at {result.path}:{warning.line} {warning.message}")
+        for warning in result.warnings:
+            click.echo(f"Warning at {result.path}:{warning.line} {warning.message}")
 
-            if result.content != result.output:
-                click.echo(f"Would sort {result.path}")
-                return_code |= 2
+        if result.content != result.output:
+            click.echo(f"Would sort {result.path}")
+            return_code |= 2
 
-            print_benchmark(result.timings)
+        print_benchmark(result.timings)
 
     return return_code
 
@@ -148,28 +147,27 @@ def diff(ctx: click.Context, filenames: List[str]) -> int:
         raise click.ClickException("Provide some filenames")
 
     return_code = 0
-    for f in filenames:
-        path = Path(f)
-        for result in usort_path(path, write=False):
-            if result.error:
-                click.echo(f"Error sorting {result.path}: {result.error}")
-                if ctx.obj.debug:
-                    click.echo(result.trace)
-                return_code |= 1
-                continue
-
-            for warning in result.warnings:
-                click.echo(f"Warning at {result.path}:{warning.line} {warning.message}")
-
-            if result.content != result.output:
-                assert result.encoding is not None
-                echo_color_unified_diff(
-                    result.content.decode(result.encoding),
-                    result.output.decode(result.encoding),
-                    result.path.as_posix(),
-                )
-
-            print_benchmark(result.timings)
+    paths = [Path(f) for f in filenames]
+    for result in usort_path(paths, write=False):
+        if result.error:
+            click.echo(f"Error sorting {result.path}: {result.error}")
+            if ctx.obj.debug:
+                click.echo(result.trace)
+            return_code |= 1
+            continue
+
+        for warning in result.warnings:
+            click.echo(f"Warning at {result.path}:{warning.line} {warning.message}")
+
+        if result.content != result.output:
+            assert result.encoding is not None
+            echo_color_unified_diff(
+                result.content.decode(result.encoding),
+                result.output.decode(result.encoding),
+                result.path.as_posix(),
+            )
+
+        print_benchmark(result.timings)
 
     return return_code
 
@@ -193,21 +191,20 @@ def format(filenames: List[str]) -> int:
         return 0 if success else 1
 
     return_code = 0
-    for f in filenames:
-        path = Path(f)
-        for result in usort_path(path, write=True):
-            if result.error:
-                click.echo(f"Error sorting {result.path}: {result.error}")
-                return_code |= 1
-                continue
+    paths = [Path(f) for f in filenames]
+    for result in usort_path(paths, write=True):
+        if result.error:
+            click.echo(f"Error sorting {result.path}: {result.error}")
+            return_code |= 1
+            continue
 
-            for warning in result.warnings:
-                click.echo(f"Warning at {result.path}:{warning.line} {warning.message}")
+        for warning in result.warnings:
+            click.echo(f"Warning at {result.path}:{warning.line} {warning.message}")
 
-            if result.content != result.output:
-                click.echo(f"Sorted {result.path}")
+        if result.content != result.output:
+            click.echo(f"Sorted {result.path}")
 
-            print_benchmark(result.timings)
+        print_benchmark(result.timings)
 
     return return_code
 


Test patch:
diff --git a/usort/tests/cli.py b/usort/tests/cli.py
index 1c41b84..fce9d01 100644
--- a/usort/tests/cli.py
+++ b/usort/tests/cli.py
@@ -55,7 +55,7 @@ class CliTest(unittest.TestCase):
                 walking \.:\s+\d+ µs
                 parsing sample\.py:\s+\d+ µs
                 sorting sample\.py:\s+\d+ µs
-                total for \.:\s+\d+ µs
+                total:\s+\d+ µs
                 """
             ).strip(),
         )
@@ -95,7 +95,7 @@ class CliTest(unittest.TestCase):
             dedent(
                 r"""
                 walking \.:\s+\d+ µs
-                total for \.:\s+\d+ µs
+                total:\s+\d+ µs
                 """
             ).strip(),
         )