Add --projects cli flag to black-primer (#2555)

* Add --projects cli flag to black-primer

Makes it possible to run a subset of projects on black primer

* Refactor into click callback

CHANGES.md

@@ -8,6 +8,7 @@
 - Add new `--workers` parameter (#2514)
 - Fixed feature detection for positional-only arguments in lambdas (#2532)
 - Bumped typed-ast version minimum to 1.4.3 for 3.10 compatiblity (#2519)
+- Add primer support for --projects (#2555)
 
 ### _Blackd_
 

mypy.ini

@@ -35,3 +35,7 @@ cache_dir=/dev/null
 [mypy-black_primer.*]
 # Until we're not supporting 3.6 primer needs this
 disallow_any_generics=False
+
+[mypy-tests.test_primer]
+# Until we're not supporting 3.6 primer needs this
+disallow_any_generics=False

src/black_primer/cli.py

@@ -1,13 +1,14 @@
 # coding=utf8
 
 import asyncio
+import json
 import logging
 import sys
 from datetime import datetime
 from pathlib import Path
 from shutil import rmtree, which
 from tempfile import gettempdir
-from typing import Any, Union, Optional
+from typing import Any, List, Optional, Union
 
 import click
 
@@ -42,12 +43,42 @@ def _handle_debug(
     return debug
 
 
+def load_projects(config_path: Path) -> List[str]:
+    with open(config_path) as config:
+        return sorted(json.load(config)["projects"].keys())
+
+
+# Unfortunately does import time file IO - but appears to be the only
+# way to get `black-primer --help` to show projects list
+DEFAULT_PROJECTS = load_projects(DEFAULT_CONFIG)
+
+
+def _projects_callback(
+    ctx: click.core.Context,
+    param: Optional[Union[click.core.Option, click.core.Parameter]],
+    projects: str,
+) -> List[str]:
+    requested_projects = set(projects.split(","))
+    available_projects = set(
+        DEFAULT_PROJECTS
+        if str(DEFAULT_CONFIG) == ctx.params["config"]
+        else load_projects(ctx.params["config"])
+    )
+
+    unavailable = requested_projects - available_projects
+    if unavailable:
+        LOG.error(f"Projects not found: {unavailable}. Available: {available_projects}")
+
+    return sorted(requested_projects & available_projects)
+
+
 async def async_main(
     config: str,
     debug: bool,
     keep: bool,
     long_checkouts: bool,
     no_diff: bool,
+    projects: List[str],
     rebase: bool,
     workdir: str,
     workers: int,
@@ -66,6 +97,7 @@ async def async_main(
             config,
             work_path,
             workers,
+            projects,
             keep,
             long_checkouts,
             rebase,
@@ -88,6 +120,8 @@ async def async_main(
     type=click.Path(exists=True),
     show_default=True,
     help="JSON config file path",
+    # Eager - because config path is used by other callback options
+    is_eager=True,
 )
 @click.option(
     "--debug",
@@ -116,6 +150,13 @@ async def async_main(
     show_default=True,
     help="Disable showing source file changes in black output",
 )
+@click.option(
+    "--projects",
+    default=",".join(DEFAULT_PROJECTS),
+    callback=_projects_callback,
+    show_default=True,
+    help="Comma separated list of projects to run",
+)
 @click.option(
     "-R",
     "--rebase",

src/black_primer/lib.py

@@ -283,16 +283,16 @@ def handle_PermissionError(
 
 async def load_projects_queue(
     config_path: Path,
+    projects_to_run: List[str],
 ) -> Tuple[Dict[str, Any], asyncio.Queue]:
     """Load project config and fill queue with all the project names"""
     with config_path.open("r") as cfp:
         config = json.load(cfp)
 
     # TODO: Offer more options here
-    # e.g. Run on X random packages or specific sub list etc.
+    # e.g. Run on X random packages etc.
-    project_names = sorted(config["projects"].keys())
+    queue: asyncio.Queue = asyncio.Queue(maxsize=len(projects_to_run))
-    queue: asyncio.Queue = asyncio.Queue(maxsize=len(project_names))
+    for project in projects_to_run:
-    for project in project_names:
         await queue.put(project)
 
     return config, queue
@@ -365,6 +365,7 @@ async def process_queue(
     config_file: str,
     work_path: Path,
     workers: int,
+    projects_to_run: List[str],
     keep: bool = False,
     long_checkouts: bool = False,
     rebase: bool = False,
@@ -383,7 +384,7 @@ async def process_queue(
     results.stats["success"] = 0
     results.stats["wrong_py_ver"] = 0
 
-    config, queue = await load_projects_queue(Path(config_file))
+    config, queue = await load_projects_queue(Path(config_file), projects_to_run)
     project_count = queue.qsize()
     s = "" if project_count == 1 else "s"
     LOG.info(f"{project_count} project{s} to run Black over")

tests/test_format.py

@@ -93,6 +93,8 @@
     "src/black/strings.py",
     "src/black/trans.py",
     "src/blackd/__init__.py",
+    "src/black_primer/cli.py",
+    "src/black_primer/lib.py",
     "src/blib2to3/pygram.py",
     "src/blib2to3/pytree.py",
     "src/blib2to3/pgen2/conv.py",

tests/test_primer.py

@@ -11,7 +11,7 @@
 from platform import system
 from subprocess import CalledProcessError
 from tempfile import TemporaryDirectory, gettempdir
-from typing import Any, Callable, Iterator, Tuple
+from typing import Any, Callable, Iterator, List, Tuple, TypeVar
 from unittest.mock import Mock, patch
 
 from click.testing import CliRunner
@@ -89,6 +89,24 @@ async def return_zero(*args: Any, **kwargs: Any) -> int:
     return 0
 
 
+if sys.version_info >= (3, 9):
+    T = TypeVar("T")
+    Q = asyncio.Queue[T]
+else:
+    T = Any
+    Q = asyncio.Queue
+
+
+def collect(queue: Q) -> List[T]:
+    ret = []
+    while True:
+        try:
+            item = queue.get_nowait()
+            ret.append(item)
+        except asyncio.QueueEmpty:
+            return ret
+
+
 class PrimerLibTests(unittest.TestCase):
     def test_analyze_results(self) -> None:
         fake_results = lib.Results(
@@ -198,10 +216,25 @@ def test_process_queue(self, mock_stdout: Mock) -> None:
         with patch("black_primer.lib.git_checkout_or_rebase", return_false):
             with TemporaryDirectory() as td:
                 return_val = loop.run_until_complete(
-                    lib.process_queue(str(config_path), Path(td), 2)
+                    lib.process_queue(
+                        str(config_path), Path(td), 2, ["django", "pyramid"]
+                    )
                 )
                 self.assertEqual(0, return_val)
 
+    @event_loop()
+    def test_load_projects_queue(self) -> None:
+        """Test the process queue on primer itself
+        - If you have non black conforming formatting in primer itself this can fail"""
+        loop = asyncio.get_event_loop()
+        config_path = Path(lib.__file__).parent / "primer.json"
+
+        config, projects_queue = loop.run_until_complete(
+            lib.load_projects_queue(config_path, ["django", "pyramid"])
+        )
+        projects = collect(projects_queue)
+        self.assertEqual(projects, ["django", "pyramid"])
+
 
 class PrimerCLITests(unittest.TestCase):
     @event_loop()
@@ -217,6 +250,7 @@ def test_async_main(self) -> None:
             "workdir": str(work_dir),
             "workers": 69,
             "no_diff": False,
+            "projects": "",
         }
         with patch("black_primer.cli.lib.process_queue", return_zero):
             return_val = loop.run_until_complete(cli.async_main(**args))  # type: ignore
@@ -230,6 +264,19 @@ def test_help_output(self) -> None:
         result = runner.invoke(cli.main, ["--help"])
         self.assertEqual(result.exit_code, 0)
 
+    def test_projects(self) -> None:
+        runner = CliRunner()
+        with event_loop():
+            result = runner.invoke(cli.main, ["--projects=tox,asdf"])
+            self.assertEqual(result.exit_code, 0)
+            assert "1 / 1 succeeded" in result.output
+
+        with event_loop():
+            runner = CliRunner()
+            result = runner.invoke(cli.main, ["--projects=tox,attrs"])
+            self.assertEqual(result.exit_code, 0)
+            assert "2 / 2 succeeded" in result.output
+
 
 if __name__ == "__main__":
     unittest.main()