summaryrefslogtreecommitdiff
path: root/doc/source/ext/snapshotqt_directive.py
blob: 1cd8dbecd92edd226a82b51a4e8c027a82e78498 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
"""RST directive to include snapshot of a Qt application in Sphinx doc.

Configuration variable in conf.py:

- snapshotqt_image_type: image file extension (default 'png').
- snapshotqt_script_dir: relative path of the root directory for scripts from
  the documentation source directory (i.e., the directory of conf.py)
  (default: '..').
"""

from __future__ import absolute_import

import logging
import os
import subprocess
import sys

from docutils.parsers.rst.directives.images import Figure


logging.basicConfig()
_logger = logging.getLogger(__name__)

# TODO:
# - Create image in build directory
# - Check if it is needed to patch block_text?

# RST directive ###############################################################

class SnapshotQtDirective(Figure):
    """Figure of a Qt application snapshot.

    Directive Type: "snapshotqt"
    Doctree Elements: As for figure
    Directive Arguments: One or more, required (script URI + script arguments).
    Directive Options: Possible.
    Directive Content: Interpreted as the figure caption and optional legend.

    A "snapshotqt" is a rst `figure
    <http://docutils.sourceforge.net/docs/ref/rst/directives.html#figure>`_
    that is generated from a Python script that uses Qt.

    The path of the script to take a snapshot is relative to
    the path given in conf.py 'snapshotqt_script_dir' value.

    ::

        .. snapshotqt: ../examples/demo.py
           :align: center
           :height: 5cm

           Caption of the figure.
    """

    # TODO this should be configured in conf.py
    IMAGE_RELATIVE_DIR = os.path.join('images')
    """The path where to store images relative to doc directory."""

    def run(self):
        # Run script stored in arguments and replace by snapshot filename

        env = self.state.document.settings.env

        # Create an image filename from arguments
        image_ext = env.config.snapshotqt_image_type.lower()
        image_name = '_'.join(self.arguments) + '.' + image_ext
        image_name = image_name.replace('./\\', '_')
        image_name = ''.join([c for c in image_name
                              if c.isalnum() or c in '_-.'])
        image_name = os.path.join(self.IMAGE_RELATIVE_DIR, image_name)

        # Change path to absolute path to run the script
        script_dir = os.path.join(env.srcdir, env.config.snapshotqt_script_dir)
        script_cmd = self.arguments[:]
        script_cmd[0] = os.path.join(script_dir, script_cmd[0])

        # Run snapshot
        snapshot_tool = os.path.abspath(__file__)
        _logger.info('Running script: %s', script_cmd)
        _logger.info('Saving snapshot to: %s', image_name)
        abs_image_name = os.path.join(env.srcdir, image_name)
        cmd = [sys.executable, snapshot_tool, '--output', abs_image_name]
        cmd += script_cmd
        subprocess.check_call(cmd)

        # Use created image as in Figure
        self.arguments = [image_name]

        return super(SnapshotQtDirective, self).run()


def setup(app):
    app.add_config_value('snapshotqt_image_type', 'png', 'env')
    app.add_config_value('snapshotqt_script_dir', '..', 'env')

    app.add_directive('snapshotqt', SnapshotQtDirective)

    return {'version': '0.1'}


# Qt monkey-patch #############################################################

def monkeyPatchQApplication(filename, qt=None):
    """Monkey-patch QApplication to take a snapshot and close the application.

    :param str filename: The image filename where to save the snapshot.
    :param str qt: The Qt binding to patch.
                   This MUST be the same as the one used by the script
                   for which to take a snapshot.
                   In: 'PyQt4', 'PyQt5', 'PySide' or None (the default).
                   If None, it will try to use PyQt4, then PySide and
                   finally PyQt4.
    """

    # Probe Qt binding
    if qt is None:
        try:
            import PyQt4.QtCore  # noqa
            qt = 'PyQt4'
        except ImportError:
            try:
                import PySide.QtCore  # noqa
                qt = 'PySide'
            except ImportError:
                try:
                    import PyQt5.QtCore  # noqa
                    qt = 'PyQt5'
                except ImportError:
                    raise RuntimeError('Cannot find any supported Qt binding.')

    if qt == 'PyQt4':
        from PyQt4.QtGui import QApplication, QPixmap
        from PyQt4.QtCore import QTimer
        import PyQt4.QtGui as _QApplicationPackage

        def grabWindow(winID):
            return QPixmap.grabWindow(winID)

    elif qt == 'PyQt5':
        from PyQt5.QtGui import QPixmap
        from PyQt5.QtWidgets import QApplication
        from PyQt5.QtCore import QTimer
        import PyQt5.QtWidgets as _QApplicationPackage

        def grabWindow(winID):
            screen = QApplication.primaryScreen()
            return screen.grabWindow(winID)

    elif qt == 'PySide':
        from PySide.QtGui import QApplication, QPixmap
        from PySide.QtCore import QTimer
        import PySide.QtGui as _QApplicationPackage

        def grabWindow(winID):
            return QPixmap.grabWindow(winID)

    else:
        raise ValueError('Unsupported Qt binding: %s' % qt)

    _logger.info('Using Qt bindings: %s', qt)

    class _QApplication(QApplication):

        _TIMEOUT = 1000.
        _FILENAME = filename

        def _grabActiveWindowAndClose(self):
            activeWindow = QApplication.activeWindow()
            if activeWindow is not None:
                if activeWindow.isVisible():
                    pixmap = grabWindow(activeWindow.winId())
                    saveOK = pixmap.save(self._FILENAME)
                    if not saveOK:
                        _logger.error(
                            'Cannot save snapshot to %s', self._FILENAME)
                else:
                    _logger.error('activeWindow is not visible.')
                self.quit()
            else:
                self._count -= 1
                if self._count > 0:
                    # Only restart a timer if everything is OK
                    QTimer.singleShot(self._TIMEOUT,
                                      self._grabActiveWindowAndClose)
                else:
                    raise RuntimeError(
                        'Aborted: It took too long to have an active window.')

        def exec_(self):
            self._count = 10
            QTimer.singleShot(self._TIMEOUT, self._grabActiveWindowAndClose)

            return super(_QApplication, self).exec_()

    _QApplicationPackage.QApplication = _QApplication


# main ########################################################################

if __name__ == '__main__':
    import argparse
    import runpy

    # Parse argv

    parser = argparse.ArgumentParser(
        description=__doc__,
        epilog="""Arguments provided after the script or module name are passed
        to the script or module.""")
    parser.add_argument(
        '-o', '--output', nargs=1, type=str,
        default='snapshot.png',
        help='Image filename of the snapshot (default: snapshot.png).')
    parser.add_argument(
        '--bindings', nargs='?',
        choices=('PySide', 'PyQt4', 'PyQt5', 'auto'),
        default='auto',
        help="""Qt bindings used by the script/module.
        If 'auto' (the default), it is probed from available python modules.
        """)
    parser.add_argument(
        '-m', '--module', action='store_true',
        help='Run the corresponding module as a script.')
    parser.add_argument(
        'script_or_module', nargs=1, type=str,
        help='Python script to run for the snapshot.')
    args, unknown = parser.parse_known_args()

    script_or_module = args.script_or_module[0]

    # arguments provided after the script or module
    extra_args = sys.argv[sys.argv.index(script_or_module) + 1:]

    if unknown != extra_args:
        parser.print_usage()
        _logger.error(
            '%s: incorrect arguments', os.path.basename(sys.argv[0]))
        sys.exit(1)

    # Monkey-patch Qt
    monkeyPatchQApplication(args.output[0],
                            args.bindings if args.bindings != 'auto' else None)

    # Update sys.argv and sys.path
    sys.argv = [script_or_module] + extra_args
    sys.path.insert(0, os.path.abspath(os.path.dirname(script_or_module)))

    if args.module:
        _logger.info('Running module: %s', ' '.join(sys.argv))
        runpy.run_module(script_or_module, run_name='__main__')

    else:
        with open(script_or_module) as f:
            code = f.read()

        _logger.info('Running script: %s', ' '.join(sys.argv))
        exec(code)