Source code for schrodinger.test.get_widget_image

#Name: get_widget_image.py
#Command: pythonrun get_widget_image.save_maestro_workspace filename
"""
Saves an image of the current workspace taken directly fro the Qt
Widget framebuffer rendering.

This can be called from Maestro as:

pythonrun get_widget_image.save_maestro_workspace filename

Copyright Schrodinger, LLC. All rights reserved.
"""

from OpenGL import GL

import schrodinger
from schrodinger.Qt.QtOpenGL import QGLWidget

maestro = schrodinger.get_maestro()


def get_qglwidget_image(widget_name=""):
    """
    Returns a QImage of a QGLWidget with optional name widget_name.
    """
    maestro.redraw()
    GL.glFlush()
    widget = maestro.get_main_window().findChild(QGLWidget, widget_name)
    if not widget:
        raise Exception('Could not locate workspace widget.')
    fbuff = widget.grabFrameBuffer()
    widget.render(fbuff)
    return fbuff


def save_maestro_workspace(filename):
    """
    Retrieves a rendering of the current Maestro workspace and saves it to
    the specified filename.
    """
    ws_image = get_qglwidget_image("main_drawing")
    ws_image.save(filename)