From 274cb2748224dd4a37e4d3ae1c6a18bd06fa3894 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Tue, 23 Jun 2020 12:57:07 -0700
Subject: [PATCH] Autodetect number of jobs with --jobs auto (#605)

This adds the possibility to autodetect the number of jobs based on the number of available CPU cores on the host system using `pytype --jobs auto`. This allows for a usage similar to [`pytest-xdist`](https://github.com/pytest-dev/pytest-xdist#speed-up-test-runs-by-sending-tests-to-multiple-cpus) which can be useful when using `pytype` in scripts that might run on different systems.

Resolves #605

PiperOrigin-RevId: 317924811
---
 pytype/tools/analyze_project/config.py     |  3 ++-
 pytype/tools/analyze_project/parse_args.py | 22 +++++++++++++++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/pytype/tools/analyze_project/config.py b/pytype/tools/analyze_project/config.py
index c97345579..4b34dec0d 100644
--- a/pytype/tools/analyze_project/config.py
+++ b/pytype/tools/analyze_project/config.py
@@ -48,7 +48,8 @@
         'Keep going past errors to analyze as many files as possible.'),
     'jobs': Item(
         1, '4', None,
-        'Run N jobs in parallel.'),
+        "Run N jobs in parallel. When 'auto' is used, this will be equivalent "
+        'to the number of CPUs on the host system.'),
     'output': Item(
         '.pytype', '.pytype', None, 'All pytype output goes here.'),
     'pythonpath': Item(
diff --git a/pytype/tools/analyze_project/parse_args.py b/pytype/tools/analyze_project/parse_args.py
index e010dd87b..1cc1a8575 100644
--- a/pytype/tools/analyze_project/parse_args.py
+++ b/pytype/tools/analyze_project/parse_args.py
@@ -1,6 +1,7 @@
 """Argument parsing for analyze_project."""
 
 import argparse
+import os
 
 from pytype import config as pytype_config
 from pytype.tools import arg_parser
@@ -10,6 +11,24 @@
 _ARG_PREFIX = '--'
 
 
+def _auto_detect_cpus():
+  try:
+    return len(os.sched_getaffinity(0))
+  except AttributeError:
+    return os.cpu_count()
+
+
+def parse_jobs(s):
+  """Parse the --jobs option."""
+  if s == 'auto':
+    n = _auto_detect_cpus()
+    return n if n else 1
+  elif s is not None:
+    return int(s)
+  else:
+    return None
+
+
 class Parser(arg_parser.Parser):
   """Subclasses Parser to add a config file processor."""
 
@@ -93,7 +112,8 @@ def make_parser():
       (('-x', '--exclude'), {'nargs': '*', 'action': 'flatten'}),
       (('inputs',), {'metavar': 'input', 'nargs': '*', 'action': 'flatten'}),
       (('-k', '--keep-going'), {'action': 'store_true', 'type': None}),
-      (('-j', '--jobs'), {'action': 'store', 'type': int, 'metavar': 'N'}),
+      (('-j', '--jobs'), {'action': 'store', 'type': parse_jobs,
+                          'metavar': 'N'}),
       (('-P', '--pythonpath'),),
       (('-V', '--python-version'),)
   ]: