utils/screenshot

#!/usr/bin/env python
# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Connect to Chameleond and capture screenshot."""

from __future__ import print_function

import argparse
import errno
import subprocess
import sys
import tempfile
import time
import xmlrpclib


def main():
  """The Main program, capture screenshot."""
  parser = argparse.ArgumentParser(
      description='Connect to Chameleond and capture screenshot.',
      formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  parser.add_argument('--chameleon_host', type=str, dest='host',
                      help='host address of Chameleond')
  parser.add_argument('--port', type=int, dest='port', default=9992,
                      help='port number of Chameleond')
  parser.add_argument('--replug', dest='replug', action='store_true',
                      help='unplug and plug before capturing screen')
  parser.add_argument('--area', type=str, dest='area',
                      help='only capture the area, e.g. \'x,y,width,height\'')
  parser.add_argument('--output', type=str, dest='output', default='image.png',
                      help='output file name of screenshot')

  options = parser.parse_args()
  if options.host is None:
    print('No chameleon_host is specified.')
    parser.print_help()
    sys.exit(errno.EINVAL)
  area_args = []
  if options.area:
    area_args = [int(s) for s in options.area.split(',')]
    if len(area_args) != 4:
      print('Invalid area parameter.')
      parser.print_help()
      sys.exit(errno.EINVAL)

  chameleon = xmlrpclib.ServerProxy(
      'http://%s:%d' % (options.host, options.port))
  inputs = chameleon.ProbeInputs()
  main_input = inputs[0]
  print('Use the main port:', chameleon.GetConnectorType(main_input))
  if options.replug:
    print('Replugging...')
    chameleon.FireHpdPulse(main_input, 1000000)
    time.sleep(5)
  width, height = chameleon.DetectResolution(main_input)
  print('Detected screen size %dx%d' % (width, height))
  if area_args:
    size_str = '%dx%d' % (area_args[2], area_args[3])
  else:
    size_str = '%dx%d' % (width, height)
  print('Capturing the screen size %s...' % size_str)
  pixels = chameleon.DumpPixels(main_input, *area_args).data
  print('Got pixel size %d' % len(pixels))
  with tempfile.NamedTemporaryFile(suffix='.rgb') as f:
    print('Temporarily write to file:', f.name)
    f.write(pixels)
    f.flush()
    command = ['convert', '-size', size_str, '-depth', '8']
    command += [f.name, options.output]
    subprocess.check_call(command)
    print('Outputted to file:', options.output)


if __name__ == '__main__':
  main()