diff --git a/caffeinate.py b/caffeinate.py index 2e309dfa6d..e59e06f0b0 100644 --- a/caffeinate.py +++ b/caffeinate.py @@ -2,6 +2,7 @@ # Use of this source code is governed by a BSD-style license that can be # found in the LICENSE file. +import contextlib import os import subprocess import sys @@ -26,3 +27,33 @@ def call(args, **call_kwargs): else: args = ['caffeinate'] + args return subprocess.call(args, **call_kwargs) + + +@contextlib.contextmanager +def scope(actually_caffeinate=True): + """Acts as a context manager keeping a Mac awake, unless flagged off. + + If the process is not running on a Mac, or `actually_caffeinate` is falsey, + this acts as a context manager that does nothing. The `actually_caffeinate` + flag is provided so command line flags can control the caffeinate behavior + without requiring weird plumbing to use or not use the context manager. + + If running on a Mac while actually_caffeinate is True (the default), this + runs `caffeinate` in a separate process, which is terminated when the + context manager exits. + """ + if sys.platform != 'darwin' or not actually_caffeinate: + # Behave like a no-op context manager. + yield False + return + + cmd = ['caffeinate', '-i', '-w', str(os.getpid())] + + proc = subprocess.Popen(cmd, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL) + try: + yield True + finally: + proc.terminate() diff --git a/fetch.py b/fetch.py index 99ba59d16a..9614b65527 100755 --- a/fetch.py +++ b/fetch.py @@ -24,6 +24,7 @@ import shlex import subprocess import sys +import caffeinate import gclient_utils import git_common @@ -215,6 +216,13 @@ def handle_args(argv): type=str, default=None, help='Protocol to use to fetch dependencies, defaults to https.') + parser.add_argument( + '--caffeinate', + action=argparse.BooleanOptionalAction, + default=True, + help=('On macOS, prevent idle sleep during the operation. Enabled by' + ' default. Use --no-caffeinate to disable. No effect on other' + ' platforms.')) parser.add_argument('config', type=str, @@ -306,8 +314,9 @@ def run(options, spec, root): def main(): args = handle_args(sys.argv) - spec, root = run_config_fetch(args.config, args.props) - return run(args, spec, root) + with caffeinate.scope(actually_caffeinate=args.caffeinate): + spec, root = run_config_fetch(args.config, args.props) + return run(args, spec, root) if __name__ == '__main__': diff --git a/tests/fetch_test.py b/tests/fetch_test.py index 5790487714..bf9277a865 100755 --- a/tests/fetch_test.py +++ b/tests/fetch_test.py @@ -49,6 +49,7 @@ class TestUtilityFunctions(unittest.TestCase): force=False, config='foo', protocol_override=None, + caffeinate=True, props=[]), response) response = fetch.handle_args([ @@ -63,11 +64,13 @@ class TestUtilityFunctions(unittest.TestCase): force=True, config='foo', protocol_override='sso', + caffeinate=True, props=['--some-param=1', '--bar=2']), response) response = fetch.handle_args([ 'filename', '-n', '--dry-run', '--no-hooks', '--nohistory', - '--force', '-p', 'sso', 'foo', '--some-param=1', '--bar=2' + '--force', '--no-caffeinate', '-p', 'sso', 'foo', '--some-param=1', + '--bar=2' ]) self.assertEqual( argparse.Namespace(dry_run=True, @@ -76,6 +79,7 @@ class TestUtilityFunctions(unittest.TestCase): force=True, config='foo', protocol_override='sso', + caffeinate=False, props=['--some-param=1', '--bar=2']), response) @mock.patch('os.path.exists', return_value=False)