Source code for screen_grab.screen_grab

# Copyright (c) 2018 Shotgun Software Inc.
# This work is provided "AS IS" and subject to the Shotgun Pipeline Toolkit
# Source Code License included in this distribution package. See LICENSE.
# By accessing, using, copying or modifying this work you indicate your
# agreement to the Shotgun Pipeline Toolkit Source Code License. All rights
# not expressly granted therein are reserved by Shotgun Software Inc.

import tempfile
import sys
import os

from sgtk.platform.qt import QtCore, QtGui
import sgtk

class ScreenGrabber(QtGui.QDialog):
    A transparent tool dialog for selecting an area (QRect) on the screen.

    This tool does not by itself perform a screen capture. The resulting
    capture rect can be used (e.g. with the get_desktop_pixmap function) to
    blit the selected portion of the screen into a pixmap.

    # If set to a callable, it will be used when performing a
    # screen grab in place of the default behavior defined in
    # this module.

    def __init__(self, parent=None):
        super(ScreenGrabber, self).__init__(parent)

        self._opacity = 1
        self._click_pos = None
        self._capture_rect = QtCore.QRect()

            | QtCore.Qt.WindowStaysOnTopHint
            | QtCore.Qt.CustomizeWindowHint
            | QtCore.Qt.Tool

        desktop = QtGui.QApplication.desktop()

    def capture_rect(self):
        The resulting QRect from a previous capture operation.
        return self._capture_rect

    def paintEvent(self, event):
        Paint event
        # Convert click and current mouse positions to local space.
        mouse_pos = self.mapFromGlobal(QtGui.QCursor.pos())
        click_pos = None
        if self._click_pos is not None:
            click_pos = self.mapFromGlobal(self._click_pos)

        painter = QtGui.QPainter(self)

        # Draw background. Aside from aesthetics, this makes the full
        # tool region accept mouse events.
        painter.setBrush(QtGui.QColor(0, 0, 0, self._opacity))

        # Clear the capture area
        if click_pos is not None:
            capture_rect = QtCore.QRect(click_pos, mouse_pos)

        pen = QtGui.QPen(QtGui.QColor(255, 255, 255, 64), 1, QtCore.Qt.DotLine)

        # Draw cropping markers at click position
        if click_pos is not None:
                event.rect().left(), click_pos.y(), event.rect().right(), click_pos.y()
                click_pos.x(), event.rect().top(), click_pos.x(), event.rect().bottom()

        # Draw cropping markers at current mouse position
            event.rect().left(), mouse_pos.y(), event.rect().right(), mouse_pos.y()
            mouse_pos.x(), event.rect().top(), mouse_pos.x(), event.rect().bottom()

    def keyPressEvent(self, event):
        Key press event
        # for some reason I am not totally sure about, it looks like
        # pressing escape while this dialog is active crashes Maya.
        # I tried subclassing closeEvent, but it looks like the crashing
        # is triggered before the code reaches this point.
        # by sealing the keypress event and not allowing any further processing
        # of the escape key (or any other key for that matter), the
        # behaviour can be successfully avoided.

        # TODO: See if we can get the behacior with hitting escape back
        # maybe by manually handling the closing of the window? I tried
        # some obvious things and weren't successful, but didn't dig very
        # deep as it felt like a nice-to-have and not a massive priority.


    def mousePressEvent(self, event):
        Mouse click event
        if event.button() == QtCore.Qt.LeftButton:
            # Begin click drag operation
            self._click_pos = event.globalPos()

    def mouseReleaseEvent(self, event):
        Mouse release event
        if event.button() == QtCore.Qt.LeftButton and self._click_pos is not None:
            # End click drag operation and commit the current capture rect
            self._capture_rect = QtCore.QRect(
                self._click_pos, event.globalPos()
            self._click_pos = None

    def mouseMoveEvent(self, event):
        Mouse move event

    def screen_capture(cls):
        Modally displays the screen capture tool.

        :returns: Captured screen
        :rtype: :class:`~PySide.QtGui.QPixmap`
        bundle = sgtk.platform.current_bundle()

        if cls.SCREEN_GRAB_CALLBACK:
            # use an external callback for screen grabbing
            return cls.SCREEN_GRAB_CALLBACK()

        elif sgtk.util.is_linux():
            # there are known issues with the QT based screen grabbing
            # on linux - some distros don't have a X11 compositing manager
            # so transparent windows aren't supported. In
            # these cases, fall back onto a traditional approach where
            # an external application is used to grab the screenshot.
            # if the external application does not exist,
            # try using the QT based approach as a fallback.
            # by using import first, we can advise users who have issues
            # with the qt approach to simply install imagemagick and things
            # should start to work.
            pixmap = _external_screenshot()

            if pixmap is None or pixmap.isNull():
                bundle.log_debug("Falling back on internal screen grabber.")
                tool = ScreenGrabber()
                pixmap = get_desktop_pixmap(tool.capture_rect)

            return pixmap

        elif sgtk.util.is_macos():
            # With macosx there are known issues with some
            # multi-diplay setups, so better to use built-in tool
            return _external_screenshot()

            # on windows, just use the QT solution.
            tool = ScreenGrabber()
            return get_desktop_pixmap(tool.capture_rect)

    def showEvent(self, event):
        Show event
        # Start fade in animation
        fade_anim = QtCore.QPropertyAnimation(self, b"_opacity_anim_prop", self)

    def _set_opacity(self, value):
        Animation callback for opacity
        self._opacity = value

    def _get_opacity(self):
        Animation callback for opacity
        return self._opacity

    _opacity_anim_prop = QtCore.Property(int, _get_opacity, _set_opacity)

    def _fit_screen_geometry(self):
        # Compute the union of all screen geometries, and resize to fit.
        desktop = QtGui.QApplication.desktop()
        workspace_rect = QtCore.QRect()
        for i in range(desktop.screenCount()):
            workspace_rect = workspace_rect.united(desktop.screenGeometry(i))

class ExternalCaptureThread(QtCore.QThread):
    Wrap external screenshot call in a thread just to be on the safe side!
    This helps avoid the os thinking the application has hung for
    certain applications (e.g. Softimage on Windows)

    def __init__(self, path):
        :param path: Path to write the screenshot to
        self._path = path
        self._error = None

    def error_message(self):
        Error message generated during capture, None if success
        return self._error

    def run(self):
            if sgtk.util.is_macos():
                # use built-in screenshot command on the mac
                ret_code = os.system("screencapture -m -i -s %s" % self._path)
                if ret_code != 0:
                    raise sgtk.TankError(
                        "Screen capture tool returned error code %s" % ret_code

            elif sgtk.util.is_linux():
                # use image magick
                ret_code = os.system("import %s" % self._path)
                if ret_code != 0:
                    raise sgtk.TankError(
                        "Screen capture tool returned error code %s. "
                        "For screen capture to work on Linux, you need to have "
                        "the imagemagick 'import' executable installed and "
                        "in your PATH." % ret_code

                raise sgtk.TankError("Unsupported platform.")
        except Exception as e:
            self._error = str(e)

def _external_screenshot():
    Use an external approach for grabbing a screenshot.
    Linux and macosx support only.

    :returns: Captured image
    :rtype: :class:`~PySide.QtGui.QPixmap`
    output_path = tempfile.NamedTemporaryFile(
        suffix=".png", prefix="screencapture_", delete=False

    pm = None
        # do screenshot with thread so we don't block anything
        screenshot_thread = ExternalCaptureThread(output_path)
        while not screenshot_thread.isFinished():

        if screenshot_thread.error_message:
            bundle = sgtk.platform.current_bundle()
                "Failed to capture " "screenshot: %s" % screenshot_thread.error_message
            # load into pixmap
            pm = QtGui.QPixmap(output_path)
        # remove the temporary file
        if output_path and os.path.exists(output_path):

    return pm

[docs]def get_desktop_pixmap(rect): """ Performs a screen capture on the specified rectangle. :param rect: Rectangle to capture :type rect: :class:`~PySide.QtCore.QRect` :returns: Captured image :rtype: :class:`~PySide.QtGui.QPixmap` """ desktop = QtGui.QApplication.desktop() return QtGui.QPixmap.grabWindow( desktop.winId(), rect.x(), rect.y(), rect.width(), rect.height() )
# Backwards compatibility, as this used to be a module-level # function but has been moved to being a classmethod on the # ScreenGrabber class. screen_capture = ScreenGrabber.screen_capture
[docs]def screen_capture_file(output_path=None): """ Modally display the screen capture tool, saving to a file. :param output_path: Path to save to. If no path is specified, a temp path is generated. :returns: path where screenshot was saved. """ if output_path is None: output_path = tempfile.NamedTemporaryFile( suffix=".png", prefix="screencapture_", delete=False ).name pixmap = screen_capture() return output_path