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 absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import codecs
import errno
from functools import reduce
import operator
import os
import subprocess
import sys
import tempfile
import time
import six
import six.moves.xmlrpc_client


class Edid(object):
    """Edid is an abstraction of Extended Display Identification Data."""

    def __init__(self, data):
        """Construct an Edid.

        @param data: A byte-array of EDID data.
        """
        self.data = data

    @classmethod
    def from_file(cls, filename):
        """Construct an Edid from a file.

        @param filename: A string of filename.
        """
        if not os.path.exists(filename):
            raise ValueError('EDID file %r does not exist' % filename)

        if filename.upper().endswith('.TXT'):
            # Convert the EDID text format, returning from xrandr.
            data = reduce(operator.add,
                          [codecs.decode(s.strip(), 'hex') for s in open(filename).readlines()])
        else:
            data = open(filename).read()
        return cls(data)


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('--edid_file', type=str, dest='edid_file', default=None,
                      help='filename of the edid to apply')
  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 = six.moves.xmlrpc_client.ServerProxy(
      'http://%s:%d' % (options.host, options.port))
  inputs = chameleon.ProbeInputs()
  main_input = inputs[0]
  print('Use the main port:', chameleon.GetConnectorType(main_input))

  was_plugged = chameleon.IsPlugged(main_input)

  if was_plugged and (options.replug or options.edid_file):
    print('Unplugging...')
    chameleon.Unplug(main_input)

  if options.edid_file:
    print('Reading EDID from %s...' % options.edid_file)
    edid = Edid.from_file(options.edid_file)
    print('Applying EDID...')
    edid_id = chameleon.CreateEdid(six.moves.xmlrpc_client.Binary(edid.data))
    chameleon.ApplyEdid(main_input, edid_id)
    chameleon.DestroyEdid(edid_id)
    time.sleep(1)

  if not chameleon.IsPlugged(main_input):
    print('Plugging...')
    chameleon.Plug(main_input)
    chameleon.WaitVideoInputStable(main_input, 30)

  try:
    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)

  finally:
    if not was_plugged:
      print('Unplugging to the original state...')
      chameleon.Unplug(main_input)


if __name__ == '__main__':
  main()