Source code for autopilot.display._screenshot

# -*- Mode: Python; coding: utf-8; indent-tabs-mode: nil; tab-width: 4 -*-
#
# Autopilot Functional Test Tool
# Copyright (C) 2014 Canonical
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

"""This module contains support for capturing screenshots."""

import logging
import os
import subprocess
import time
import tempfile
from io import BytesIO

from PIL import Image

import autopilot._glib

logger = logging.getLogger(__name__)


[docs]def get_screenshot_data(display_type): """Return a BytesIO object of the png data for the screenshot image. *display_type* is the display server type. supported values are: - "X11" - "MIR" :raises RuntimeError: If attempting to capture an image on an unsupported display server. :raises RuntimeError: If saving image data to file-object fails. """ if display_type == "MIR": return _get_screenshot_mir() elif display_type == "X11": return _get_screenshot_x11() else: raise RuntimeError( "Don't know how to take screen shot for this display server: {}" .format(display_type) )
def _get_screenshot_x11(): """Capture screenshot from an X11 display. :raises RuntimeError: If saving pixbuf to fileobject fails. """ pixbuf_data = _get_x11_pixbuf_data() return _save_gdk_pixbuf_to_fileobject(pixbuf_data) def _get_x11_pixbuf_data(): Gdk = autopilot._glib._import_gdk() window = Gdk.get_default_root_window() x, y, width, height = window.get_geometry() return Gdk.pixbuf_get_from_window(window, x, y, width, height) def _save_gdk_pixbuf_to_fileobject(pixbuf): image_data = pixbuf.save_to_bufferv("png", [], []) if image_data[0] is True: image_datafile = BytesIO() image_datafile.write(image_data[1]) image_datafile.seek(0) return image_datafile logger.error("Unable to write image data.") raise RuntimeError("Failed to save image data to file object.") def _get_screenshot_mir(): """Capture screenshot from Mir display. :raises FileNotFoundError: If the mirscreencast utility is not found. :raises CalledProcessError: If the mirscreencast utility errors while taking a screenshot. :raises ValueError: If the PNG conversion step fails. """ from autopilot.display import Display display_resolution = Display.create().get_screen_geometry(0)[2:] screenshot_filepath = _take_mirscreencast_screenshot() try: png_data_file = _get_png_from_rgba_file( screenshot_filepath, display_resolution ) finally: os.remove(screenshot_filepath) return png_data_file def _take_mirscreencast_screenshot(): """Takes a single frame capture of the screen using mirscreencast. Return the path to the resulting rgba file. :raises FileNotFoundError: If the mirscreencast utility is not found. :raises CalledProcessError: If the mirscreencast utility errors while taking a screenshot. """ timestamp = int(time.time()) filename = "ap-screenshot-data-{ts}.rgba".format(ts=timestamp) filepath = os.path.join(tempfile.gettempdir(), filename) try: subprocess.check_call([ "mirscreencast", "-m", "/run/mir_socket", "-n", "1", "-f", filepath ]) except FileNotFoundError as e: e.args += ("The utility 'mirscreencast' is not available.", ) raise except subprocess.CalledProcessError as e: e.args += ("Failed to take screenshot.", ) raise return filepath # This currently uses PIL but I'm investigating using something # quicker/lighter-weight. def _get_png_from_rgba_file(filepath, image_size): """Convert an rgba file to a png file stored in a filelike object. Returns a BytesIO object containing the png data. """ image_data = _image_data_from_file(filepath, image_size) bio = BytesIO() image_data.save(bio, format="png") bio.seek(0) return bio def _image_data_from_file(filepath, image_size): with open(filepath, "rb") as f: image_data = Image.frombuffer( "RGBA", image_size, f.read(), "raw", "RGBA", 0, 1 ) return image_data