Add :screenshot command

See #1841 and #2146

6 files changed, 144 insertions, 5 deletions

diff --git a/doc/changelog.asciidoc b/doc/changelog.asciidoc
index 4307c6cf8..33631fe3b 100644
--- a/doc/changelog.asciidoc
+++ b/doc/changelog.asciidoc
@@ -19,6 +19,12 @@ breaking changes (such as renamed commands) can happen in minor releases.
 v2.1.0 (unreleased)
 -------------------
 
+Added
+~~~~~
+
+- New `:screenshot` command which can be used to screenshot the visible part of
+  the page.
+
 Changed
 ~~~~~~~
 
diff --git a/doc/help/commands.asciidoc b/doc/help/commands.asciidoc
index 014dab2df..fdd485b28 100644
--- a/doc/help/commands.asciidoc
+++ b/doc/help/commands.asciidoc
@@ -99,6 +99,7 @@ possible to run or bind multiple commands by separating them with `;;`.
 |<<restart,restart>>|Restart qutebrowser while keeping existing tabs open.
 |<<run-with-count,run-with-count>>|Run a command with the given count.
 |<<save,save>>|Save configs and state.
+|<<screenshot,screenshot>>|Take a screenshot of the currently shown part of the page.
 |<<scroll,scroll>>|Scroll the current tab in the given direction.
 |<<scroll-page,scroll-page>>|Scroll the frame page-wise.
 |<<scroll-px,scroll-px>>|Scroll the current tab by 'count * dx/dy' pixels.
@@ -1080,6 +1081,21 @@ Save configs and state.
 * +'what'+: What to save (`config`/`key-config`/`cookies`/...). If not given, everything is saved.
 
 
+[[screenshot]]
+=== screenshot
+Syntax: +:screenshot [*--rect* 'rect'] [*--force*] 'filename'+
+
+Take a screenshot of the currently shown part of the page.
+
+The file format is automatically determined based on the given file extension.
+
+==== positional arguments
+* +'filename'+: The file to save the screenshot to (~ gets expanded).
+
+==== optional arguments
+* +*-r*+, +*--rect*+: The rectangle to save, as a string like WxH+X+Y.
+* +*-f*+, +*--force*+: Overwrite existing files.
+
 [[scroll]]
 === scroll
 Syntax: +:scroll 'direction'+
diff --git a/qutebrowser/browser/browsertab.py b/qutebrowser/browser/browsertab.py
index 6b73c92b8..634adbee0 100644
--- a/qutebrowser/browser/browsertab.py
+++ b/qutebrowser/browser/browsertab.py
@@ -27,8 +27,8 @@ from typing import (cast, TYPE_CHECKING, Any, Callable, Iterable, List, Optional
                     Sequence, Set, Type, Union)
 
 from PyQt5.QtCore import (pyqtSignal, pyqtSlot, QUrl, QObject, QSizeF, Qt,
-                          QEvent, QPoint)
-from PyQt5.QtGui import QKeyEvent, QIcon
+                          QEvent, QPoint, QRect)
+from PyQt5.QtGui import QKeyEvent, QIcon, QPixmap
 from PyQt5.QtWidgets import QWidget, QApplication, QDialog
 from PyQt5.QtPrintSupport import QPrintDialog, QPrinter
 from PyQt5.QtNetwork import QNetworkAccessManager
@@ -1199,6 +1199,12 @@ class AbstractTab(QWidget):
         """
         raise NotImplementedError
 
+    def grab_pixmap(self, rect: QRect = QRect()) -> QPixmap:
+        """Grab a QPixmap of the displayed page."""
+        pic = self._widget.grab(rect)
+        assert not pic.isNull()
+        return pic
+
     def __repr__(self) -> str:
         try:
             qurl = self.url()
diff --git a/qutebrowser/components/misccommands.py b/qutebrowser/components/misccommands.py
index b74998ae1..0ead38799 100644
--- a/qutebrowser/components/misccommands.py
+++ b/qutebrowser/components/misccommands.py
@@ -23,6 +23,7 @@ import os
 import signal
 import functools
 import logging
+import pathlib
 from typing import Optional
 
 try:
@@ -30,13 +31,14 @@ try:
 except ImportError:
     hunter = None
 
-from PyQt5.QtCore import Qt
+from PyQt5.QtCore import Qt, QRect
 from PyQt5.QtPrintSupport import QPrintPreviewDialog
 
 from qutebrowser.api import cmdutils, apitypes, message, config
 
 # FIXME should be part of qutebrowser.api?
 from qutebrowser.completion.models import miscmodels
+from qutebrowser.utils import utils
 
 
 @cmdutils.register(name='reload')
@@ -155,6 +157,37 @@ def debug_dump_page(tab: apitypes.Tab, dest: str, plain: bool = False) -> None:
     tab.dump_async(callback, plain=plain)
 
 
+@cmdutils.register()
+@cmdutils.argument('tab', value=cmdutils.Value.cur_tab)
+def screenshot(
+        tab: apitypes.Tab,
+        filename: pathlib.Path,
+        *,
+        rect: str = None,
+        force: bool = False,
+) -> None:
+    """Take a screenshot of the currently shown part of the page.
+
+    The file format is automatically determined based on the given file extension.
+
+    Args:
+        filename: The file to save the screenshot to (~ gets expanded).
+        rect: The rectangle to save, as a string like WxH+X+Y.
+        force: Overwrite existing files.
+    """
+    expanded = filename.expanduser()
+    if expanded.exists() and not force:
+        raise cmdutils.CommandError(
+            f"File {filename} already exists (use --force to overwrite)")
+
+    qrect = QRect() if rect is None else utils.parse_rect(rect)
+
+    pic = tab.grab_pixmap(qrect)
+    ok = pic.save(str(expanded))
+    if not ok:
+        raise cmdutils.CommandError(f"Saving to {filename} failed")
+
+
 @cmdutils.register(maxsplit=0)
 @cmdutils.argument('tab', value=cmdutils.Value.cur_tab)
 def insert_text(tab: apitypes.Tab, text: str) -> None:
diff --git a/qutebrowser/utils/utils.py b/qutebrowser/utils/utils.py
index 515d73fc6..698a608ef 100644
--- a/qutebrowser/utils/utils.py
+++ b/qutebrowser/utils/utils.py
@@ -47,7 +47,7 @@ except ImportError:  # pragma: no cover
 
             """Empty stub at runtime."""
 
-from PyQt5.QtCore import QUrl, QVersionNumber
+from PyQt5.QtCore import QUrl, QVersionNumber, QRect
 from PyQt5.QtGui import QClipboard, QDesktopServices
 from PyQt5.QtWidgets import QApplication
 # We cannot use the stdlib version on 3.7-3.8 because we need the files() API.
@@ -906,3 +906,31 @@ def cleanup_file(filepath: str) -> Iterator[None]:
             os.remove(filepath)
         except OSError as e:
             log.misc.error(f"Failed to delete tempfile {filepath} ({e})!")
+
+
+_RECT_PATTERN = re.compile(r'(?P<w>\d+)x(?P<h>\d+)\+(?P<x>\d+)\+(?P<y>\d+)')
+
+
+def parse_rect(s: str) -> QRect:
+    """Parse a rectangle string like 20x20+5+3.
+
+    Negative offsets aren't supported, and neither is leaving off parts of the string.
+    """
+    match = _RECT_PATTERN.match(s)
+    if not match:
+        raise ValueError(f"String {s} does not match WxH+X+Y")
+
+    w = int(match.group('w'))
+    h = int(match.group('h'))
+    x = int(match.group('x'))
+    y = int(match.group('y'))
+
+    try:
+        rect = QRect(x, y, w, h)
+    except OverflowError as e:
+        raise ValueError(e)
+
+    if not rect.isValid():
+        raise ValueError("Invalid rectangle")
+
+    return rect
diff --git a/tests/unit/utils/test_utils.py b/tests/unit/utils/test_utils.py
index 2776b1f30..4cf60943c 100644
--- a/tests/unit/utils/test_utils.py
+++ b/tests/unit/utils/test_utils.py
@@ -30,7 +30,7 @@ import shlex
 import math
 import zipfile
 
-from PyQt5.QtCore import QUrl
+from PyQt5.QtCore import QUrl, QRect
 from PyQt5.QtGui import QClipboard
 import pytest
 import hypothesis
@@ -1009,3 +1009,53 @@ class TestCleanupFileContext:
                 pass
         assert len(caplog.messages) == 1
         assert caplog.messages[0].startswith("Failed to delete tempfile")
+
+
+class TestParseRect:
+
+    @pytest.mark.parametrize('value, expected', [
+        ('1x1+0+0', QRect(0, 0, 1, 1)),
+        ('123x789+12+34', QRect(12, 34, 123, 789)),
+    ])
+    def test_valid(self, value, expected):
+        assert utils.parse_rect(value) == expected
+
+    @pytest.mark.parametrize('value, message', [
+        ('0x0+1+1', "Invalid rectangle"),
+        ('1x1-1+1', "String 1x1-1+1 does not match WxH+X+Y"),
+        ('1x1+1-1', "String 1x1+1-1 does not match WxH+X+Y"),
+        ('1x1', "String 1x1 does not match WxH+X+Y"),
+        ('+1+2', "String +1+2 does not match WxH+X+Y"),
+        ('1e0x1+0+0', "String 1e0x1+0+0 does not match WxH+X+Y"),
+        ('¹x1+0+0', "String ¹x1+0+0 does not match WxH+X+Y"),
+    ])
+    def test_invalid(self, value, message):
+        with pytest.raises(ValueError) as excinfo:
+            utils.parse_rect(value)
+        assert str(excinfo.value) == message
+
+    @hypothesis.given(strategies.text())
+    def test_hypothesis_text(self, s):
+        try:
+            utils.parse_rect(s)
+        except ValueError as e:
+            print(e)
+
+    @hypothesis.given(strategies.tuples(
+        strategies.integers(),
+        strategies.integers(),
+        strategies.integers(),
+        strategies.integers(),
+    ).map(lambda tpl: '{}x{}+{}+{}'.format(*tpl)))
+    def test_hypothesis_sophisticated(self, s):
+        try:
+            utils.parse_rect(s)
+        except ValueError as e:
+            print(e)
+
+    @hypothesis.given(strategies.from_regex(utils._RECT_PATTERN))
+    def test_hypothesis_regex(self, s):
+        try:
+            utils.parse_rect(s)
+        except ValueError as e:
+            print(e)