Merge pull request facebook#161 from jreese/gather

Optimize passing multiple paths to usort CLI

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,24 +117,35 @@ 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.
 
     If given a directory, it will be searched, recursively, for any Python source files,
     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
 
 

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
 

usort/tests/cli.py

@@ -55,7 +55,7 @@ def test_benchmark(self) -> None:
                 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 @@ def test_benchmark_multiple(self) -> None:
             dedent(
                 r"""
                 walking \.:\s+\d+ µs
-                total for \.:\s+\d+ µs
+                total:\s+\d+ µs
                 """
             ).strip(),
         )