From 45573570cf706f8490cd631ee6f6a58bc41f8130 Mon Sep 17 00:00:00 2001
From: Ruairidh MacLeod <5160559+rkm@users.noreply.github.com>
Date: Tue, 12 May 2020 14:23:26 +0100
Subject: [PATCH] Parse --jobs as a custom argparse type. Fixes #567

---
 src/flake8/checker.py              | 11 ++---------
 src/flake8/main/options.py         | 22 ++++++++++++++++++++++
 tests/unit/test_checker_manager.py |  3 ++-
 tests/unit/test_option_manager.py  | 30 ++++++++++++++++++++++++++++++
 4 files changed, 56 insertions(+), 10 deletions(-)

diff --git a/src/flake8/checker.py b/src/flake8/checker.py
index 36a4735..d993cb9 100644
--- a/src/flake8/checker.py
+++ b/src/flake8/checker.py
@@ -137,19 +137,12 @@ class Manager(object):
             return 0
 
         jobs = self.options.jobs
-        if jobs != "auto" and not jobs.isdigit():
-            LOG.warning(
-                '"%s" is not a valid parameter to --jobs. Must be one '
-                'of "auto" or a numerical value, e.g., 4.',
-                jobs,
-            )
-            return 0
 
         # If the value is "auto", we want to let the multiprocessing library
         # decide the number based on the number of CPUs. However, if that
         # function is not implemented for this particular value of Python we
         # default to 1
-        if jobs == "auto":
+        if jobs.is_auto:
             try:
                 return multiprocessing.cpu_count()
             except NotImplementedError:
@@ -157,7 +150,7 @@ class Manager(object):
 
         # Otherwise, we know jobs should be an integer and we can just convert
         # it to an integer
-        return int(jobs)
+        return jobs.n_jobs
 
     def _handle_results(self, filename, results):
         style_guide = self.style_guide
diff --git a/src/flake8/main/options.py b/src/flake8/main/options.py
index ba1f1c2..f4588e8 100644
--- a/src/flake8/main/options.py
+++ b/src/flake8/main/options.py
@@ -62,6 +62,27 @@ def register_preliminary_options(parser):
     )
 
 
+class JobsArgument:
+    """Type callback for the --jobs argument."""
+
+    def __init__(self, arg):  # type: (str) -> None
+        """Parse and validate the --jobs argument.
+
+        :param str arg:
+            The argument passed by argparse for validation
+        """
+        self.is_auto = False
+        self.n_jobs = -1
+        if arg == "auto":
+            self.is_auto = True
+        elif arg.isdigit():
+            self.n_jobs = int(arg)
+        else:
+            raise argparse.ArgumentTypeError(
+                "{!r} must be 'auto' or an integer.".format(arg),
+            )
+
+
 def register_default_options(option_manager):
     """Register the default options on our OptionManager.
 
@@ -293,6 +314,7 @@ def register_default_options(option_manager):
         "--jobs",
         default="auto",
         parse_from_config=True,
+        type=JobsArgument,
         help="Number of subprocesses to use to run checks in parallel. "
         'This is ignored on Windows. The default, "auto", will '
         "auto-detect the number of processors available to use."
diff --git a/tests/unit/test_checker_manager.py b/tests/unit/test_checker_manager.py
index 1d7e547..d3e7e61 100644
--- a/tests/unit/test_checker_manager.py
+++ b/tests/unit/test_checker_manager.py
@@ -5,13 +5,14 @@ import mock
 import pytest
 
 from flake8 import checker
+from flake8.main.options import JobsArgument
 
 
 def style_guide_mock():
     """Create a mock StyleGuide object."""
     return mock.MagicMock(**{
         'options.diff': False,
-        'options.jobs': '4',
+        'options.jobs': JobsArgument("4"),
     })
 
 
diff --git a/tests/unit/test_option_manager.py b/tests/unit/test_option_manager.py
index f2ee4f9..3a63db9 100644
--- a/tests/unit/test_option_manager.py
+++ b/tests/unit/test_option_manager.py
@@ -6,6 +6,7 @@ import mock
 import pytest
 
 from flake8 import utils
+from flake8.main.options import JobsArgument
 from flake8.options import manager
 
 TEST_VERSION = '3.0.0b1'
@@ -343,3 +344,32 @@ def test_optmanager_group(optmanager, capsys):
     out, err = capsys.readouterr()
     output = out + err
     assert '\ngroupname:\n' in output
+
+
+@pytest.mark.parametrize(
+    ("s", "is_auto", "n_jobs"),
+    (
+        ("auto", True, -1),
+        ("4", False, 4),
+    ),
+)
+def test_parse_valid_jobs_argument(s, is_auto, n_jobs):
+    """Test that --jobs properly parses valid arguments."""
+    jobs_opt = JobsArgument(s)
+    assert is_auto == jobs_opt.is_auto
+    assert n_jobs == jobs_opt.n_jobs
+
+
+def test_parse_invalid_jobs_argument(optmanager, capsys):
+    """Test that --jobs properly rejects invalid arguments."""
+    namespace = argparse.Namespace()
+    optmanager.add_option("--jobs", type=JobsArgument)
+    with pytest.raises(SystemExit):
+        optmanager.parse_args(["--jobs=foo"], namespace)
+    out, err = capsys.readouterr()
+    output = out + err
+    expected = (
+        "\nflake8: error: argument --jobs: "
+        "'foo' must be 'auto' or an integer.\n"
+    )
+    assert expected in output