devil/devil/android/tools/screenshot.py

#!/usr/bin/env python
# Copyright 2015 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Takes a screenshot from an Android device."""

import argparse
import logging
import os
import sys

if __name__ == '__main__':
  sys.path.append(
      os.path.abspath(
          os.path.join(os.path.dirname(__file__), '..', '..', '..')))
from devil.android import device_utils
from devil.android.tools import script_common
from devil.utils import logging_common

logger = logging.getLogger(__name__)


def main():
  # Parse options.
  parser = argparse.ArgumentParser(description=__doc__)
  logging_common.AddLoggingArguments(parser)
  script_common.AddDeviceArguments(parser)
  parser.add_argument(
      '-f',
      '--file',
      metavar='FILE',
      help='Save result to file instead of generating a '
      'timestamped file name.')
  parser.add_argument(
      'host_file',
      nargs='?',
      help='File to which the screenshot will be saved.')

  args = parser.parse_args()
  host_file = args.host_file or args.file
  logging_common.InitializeLogging(args)

  devices = script_common.GetDevices(args.devices, args.denylist_file)

  def screenshot(device):
    f = None
    if host_file:
      root, ext = os.path.splitext(host_file)
      f = '%s_%s%s' % (root, str(device), ext)
    f = device.TakeScreenshot(f)
    print('Screenshot for device %s written to %s' %
          (str(device), os.path.abspath(f)))

  device_utils.DeviceUtils.parallel(devices).pMap(screenshot)
  return 0


if __name__ == '__main__':
  sys.exit(main())