Repository: facebookexperimental/usort
Pull request #70 (2021-10-02 00:48:07)

Use trailrunner for walk and multiprocessing
- Uses `trailrunner.walk()` to find paths to sort.
- Adds `usort_file()` function to sort a single file.
- Use `trailrunner.run()` in `usort_path()` to use child process for
  each file being sorted.
- Modify existing benchmark/timing APIs and usage to account for
  multiprocessing usage and return timings in the `Result` object.
- Update `cli` to use new timings API for printing benchmarks

Fixes #61 

jreese: I benchmarked running usort on my LibCST checkout, using my M1 Macbook.

Running usort from main:
```
$workspace/usort/.venv/bin/usort diff libcst/  12.77s user 0.16s system 97% cpu 13.312 total
$workspace/usort/.venv/bin/usort diff libcst/  12.59s user 0.13s system 97% cpu 13.116 total
$workspace/usort/.venv/bin/usort diff libcst/  12.67s user 0.12s system 97% cpu 13.123 total
```

With this PR:
```
$workspace/usort/.venv/bin/usort diff libcst/  21.32s user 0.59s system 477% cpu 4.586 total
$workspace/usort/.venv/bin/usort diff libcst/  21.29s user 0.60s system 487% cpu 4.491 total
$workspace/usort/.venv/bin/usort diff libcst/  21.35s user 0.60s system 486% cpu 4.513 total
```


Patch:
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
deleted file mode 100644
index 3afb273..0000000
--- a/.github/dependabot.yml
+++ /dev/null
@@ -1,7 +0,0 @@
-version: 2
-updates:
-  - package-ecosystem: "pip"
-    directory: "/"
-    schedule:
-      interval: "weekly"
-      day: "tuesday"
\ No newline at end of file
diff --git a/requirements-dev.txt b/requirements-dev.txt
index b3dfc6b..e3fa404 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -1,5 +1,5 @@
 black==21.9b0
-coverage==6.0
+coverage==5.5
 flake8==3.9.2
 mypy==0.910
 sphinx==4.2.0
@@ -7,5 +7,5 @@ tox==3.24.4
 twine==3.4.2
 ufmt==1.3.0
 volatile==2.1.0
-wheel==0.37.0
+wheel==0.33.6
 git+https://github.com/miyakogi/m2r@66f4a5a500cdd9fc59085106bff082c9cadafaf3
diff --git a/requirements.txt b/requirements.txt
index ffc3849..1a35ad1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,3 +2,4 @@ LibCST==0.3.21
 stdlibs==2021.4.1
 moreorless==0.4.0
 toml==0.10.2
+trailrunner==1.1.1
diff --git a/setup.cfg b/setup.cfg
index 474884d..6c25451 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -19,6 +19,7 @@ install_requires =
     moreorless >= 0.3.0
     stdlibs >= 2021.4.1
     toml >= 0.10.0
+    trailrunner >= 1.0
 
 [options.entry_points]
 console_scripts =
diff --git a/usort/__init__.py b/usort/__init__.py
index bb6f454..bf5db27 100644
--- a/usort/__init__.py
+++ b/usort/__init__.py
@@ -10,11 +10,12 @@ try:
 except ImportError:
     __version__ = "dev"
 
-from .api import usort_bytes, usort_path, usort_stdin, usort_string
+from .api import usort_bytes, usort_file, usort_path, usort_stdin, usort_string
 
 __all__ = [
     "__version__",
     "usort_bytes",
+    "usort_file",
     "usort_path",
     "usort_stdin",
     "usort_string",
diff --git a/usort/api.py b/usort/api.py
index c1feed8..0f46214 100644
--- a/usort/api.py
+++ b/usort/api.py
@@ -4,13 +4,16 @@
 # LICENSE file in the root directory of this source tree.
 
 import sys
+from functools import partial
 from pathlib import Path
 from typing import Iterable, Optional, Tuple
 
+from trailrunner import walk, run
+
 from .config import Config
 from .sorting import ImportSortingTransformer
 from .types import Result
-from .util import try_parse, timed, walk
+from .util import get_timings, try_parse, timed
 
 
 __all__ = ["usort_bytes", "usort_string", "usort_path", "usort_stdin"]
@@ -48,28 +51,40 @@ def usort_string(data: str, config: Config, path: Optional[Path] = None) -> str:
     return usort_bytes(data=data.encode(), config=config, path=path)[0].decode()
 
 
+def usort_file(path: Path, *, write: bool = False) -> Result:
+    """
+    Format a single file and return a Result object.
+    """
+
+    data: bytes = b""
+    try:
+        config = Config.find(path.parent)
+        data = path.read_bytes()
+        output, encoding = usort_bytes(data, config, path)
+        if write:
+            path.write_bytes(output)
+        return Result(
+            path=path,
+            content=data,
+            output=output,
+            encoding=encoding,
+            timings=get_timings(),
+        )
+
+    except Exception as e:
+        return Result(path=path, content=data, error=e, timings=get_timings())
+
+
 def usort_path(path: Path, *, write: bool = False) -> Iterable[Result]:
     """
     For a given path, format it, or any .py files in it, and yield Result objects
     """
-    files: Iterable[Path]
-    if path.is_dir():
-        files = walk(path, "*.py")
-    else:
-        files = [path]
+    with timed(f"total for {path}"):
+        with timed(f"walking {path}"):
+            paths = walk(path)
 
-    data: bytes = b""
-    for f in files:
-        try:
-            config = Config.find(f.parent)
-            data = f.read_bytes()
-            output, encoding = usort_bytes(data, config, f)
-            if write:
-                f.write_bytes(output)
-            yield Result(f, data, output, encoding)
-
-        except Exception as e:
-            yield Result(f, data, b"", None, e)
+        fn = partial(usort_file, write=write)
+        return (v for v in run(paths, fn).values())
 
 
 def usort_stdin() -> bool:
diff --git a/usort/cli.py b/usort/cli.py
index c0aa697..4d8575f 100644
--- a/usort/cli.py
+++ b/usort/cli.py
@@ -6,7 +6,7 @@
 import sys
 from functools import wraps
 from pathlib import Path
-from typing import Any, Callable, List
+from typing import Any, Callable, List, Sequence
 
 import click
 from moreorless.click import echo_color_unified_diff
@@ -15,11 +15,17 @@ from . import __version__
 from .api import usort_path, usort_stdin
 from .config import Config
 from .sorting import sortable_blocks
-from .util import TIMINGS, print_timings, try_parse
+from .types import Timing
+from .util import get_timings, print_timings, try_parse
 
 BENCHMARK = False
 
 
+def print_benchmark(timings: Sequence[Timing]) -> None:
+    if BENCHMARK:
+        print_timings(click.echo, timings=timings)
+
+
 def usort_command(fn: Callable[..., int]) -> Callable[..., None]:
     """
     Run wrapped command, print timings if --benchmark, and exit with return code
@@ -27,10 +33,9 @@ def usort_command(fn: Callable[..., int]) -> Callable[..., None]:
 
     @wraps(fn)
     def wrapper(*args: Any, **kwargs: Any) -> None:
-        TIMINGS.clear()
         exit_code = fn(*args, **kwargs) or 0
         if BENCHMARK:
-            print_timings(click.echo)
+            print_timings(click.echo, timings=get_timings())
         sys.exit(exit_code)
 
     return wrapper
@@ -110,6 +115,8 @@ def check(filenames: List[str]) -> int:
                 click.echo(f"Would sort {result.path}")
                 return_code |= 2
 
+            print_benchmark(result.timings)
+
     return return_code
 
 
@@ -139,6 +146,8 @@ def diff(filenames: List[str]) -> int:
                     result.path.as_posix(),
                 )
 
+            print_benchmark(result.timings)
+
     return return_code
 
 
@@ -172,6 +181,8 @@ def format(filenames: List[str]) -> int:
             if result.content != result.output:
                 click.echo(f"Sorted {result.path}")
 
+            print_benchmark(result.timings)
+
     return return_code
 
 
diff --git a/usort/types.py b/usort/types.py
index a29b72b..33d32af 100644
--- a/usort/types.py
+++ b/usort/types.py
@@ -5,22 +5,25 @@
 
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Sequence, Tuple
 
 import libcst as cst
 
 from .config import Config
 
+Timing = Tuple[str, float]
+
 
 @dataclass
 class Result:
     path: Path
     content: bytes
-    output: bytes
+    output: bytes = b""
     # encoding will be None on parse errors; we get this from LibCST on a successful
     # parse.
-    encoding: Optional[str]
+    encoding: Optional[str] = None
     error: Optional[Exception] = None
+    timings: Sequence[Timing] = ()
 
 
 @dataclass(order=True)
diff --git a/usort/util.py b/usort/util.py
index 22f5b00..afb1362 100644
--- a/usort/util.py
+++ b/usort/util.py
@@ -3,16 +3,16 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-import os
 from contextlib import contextmanager
-from fnmatch import fnmatch
 from pathlib import Path
 from time import monotonic
-from typing import Callable, Generator, Iterable, List, Optional, Tuple
+from typing import Callable, Generator, List, Optional, Tuple, Sequence
 
 import libcst as cst
 
-TIMINGS: List[Tuple[str, float]] = []
+from .types import Timing
+
+TIMINGS: List[Timing] = []
 
 
 @contextmanager
@@ -28,26 +28,23 @@ def timed(msg: str) -> Generator[None, None, None]:
     TIMINGS.append((msg, after - before))
 
 
-def print_timings(fn: Callable[[str], None] = print) -> None:
+def get_timings() -> Sequence[Tuple[str, float]]:
+    try:
+        return list(TIMINGS)
+    finally:
+        TIMINGS.clear()
+
+
+def print_timings(
+    fn: Callable[[str], None] = print, *, timings: Sequence[Timing]
+) -> None:
     """
     Print all stored timing values in microseconds.
     """
-    for msg, duration in TIMINGS:
+    for msg, duration in timings:
         fn(f"{msg + ':':50} {int(duration*1000000):7} µs")
 
 
-def walk(path: Path, glob: str) -> Iterable[Path]:
-    with timed(f"walking {path}"):
-        paths: List[Path] = []
-        for root, dirs, files in os.walk(path):
-            dirs[:] = [d for d in dirs if not d.startswith(".")]
-            root_path = Path(root)
-            for f in files:
-                if fnmatch(f, glob):
-                    paths.append(root_path / f)
-        return paths
-
-
 def try_parse(path: Path, data: Optional[bytes] = None) -> cst.Module:
     """
     Attempts to parse the file with all syntax versions known by LibCST.


Test patch:
diff --git a/usort/tests/cli.py b/usort/tests/cli.py
index f9b15df..e802249 100644
--- a/usort/tests/cli.py
+++ b/usort/tests/cli.py
@@ -7,6 +7,7 @@ import os
 import unittest
 from contextlib import contextmanager
 from pathlib import Path
+from textwrap import dedent
 from typing import AnyStr, Generator
 
 import volatile
@@ -46,10 +47,14 @@ class CliTest(unittest.TestCase):
 
         self.assertRegex(
             result.output,
-            r"""walking \.:\s+\d+ µs
-parsing sample\.py:\s+\d+ µs
-sorting sample\.py:\s+\d+ µs
-""",
+            dedent(
+                r"""
+                parsing sample\.py:\s+\d+ µs
+                sorting sample\.py:\s+\d+ µs
+                walking \.:\s+\d+ µs
+                total for \.:\s+\d+ µs
+                """
+            ).strip(),
         )
         self.assertEqual(0, result.exit_code)