From 4a9e22163b232038a0ad729fdded9c2281a45505 Mon Sep 17 00:00:00 2001
From: Ryan Roden-Corrent <ryan@rcorre.net>
Date: Mon, 2 Oct 2017 12:01:53 -0400
Subject: [PATCH] Filter keyhints based on count prefix.

If a count prefix is given, only hint commands that can take a count.
---
 qutebrowser/commands/command.py   |  4 ++++
 qutebrowser/misc/keyhintwidget.py | 11 +++++++++--
 tests/helpers/stubs.py            |  1 +
 tests/unit/misc/test_keyhints.py  | 13 +++++++++----
 4 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/qutebrowser/commands/command.py b/qutebrowser/commands/command.py
index fd9b2382f..9baa4efe7 100644
--- a/qutebrowser/commands/command.py
+++ b/qutebrowser/commands/command.py
@@ -515,3 +515,7 @@ class Command:
             raise cmdexc.PrerequisitesError(
                 "{}: This command is only allowed in {} mode, not {}.".format(
                     self.name, mode_names, mode.name))
+
+    def takes_count(self):
+        """Return true iff this command can take a count argument."""
+        return any(arg.count for arg in self._qute_args)
diff --git a/qutebrowser/misc/keyhintwidget.py b/qutebrowser/misc/keyhintwidget.py
index ad2ea84a2..5626ce4ae 100644
--- a/qutebrowser/misc/keyhintwidget.py
+++ b/qutebrowser/misc/keyhintwidget.py
@@ -33,6 +33,7 @@ from PyQt5.QtCore import pyqtSlot, pyqtSignal, Qt
 
 from qutebrowser.config import config
 from qutebrowser.utils import utils, usertypes
+from qutebrowser.commands import cmdutils
 
 
 class KeyHintView(QLabel):
@@ -86,7 +87,7 @@ class KeyHintView(QLabel):
         Args:
             prefix: The current partial keystring.
         """
-        _, prefix = re.match(r'^(\d*)(.*)', prefix).groups()
+        countstr, prefix = re.match(r'^(\d*)(.*)', prefix).groups()
         if not prefix:
             self._show_timer.stop()
             self.hide()
@@ -96,11 +97,17 @@ class KeyHintView(QLabel):
             return any(fnmatch.fnmatchcase(keychain, glob)
                        for glob in config.val.keyhint.blacklist)
 
+        def takes_count(cmdstr):
+            cmdname = cmdstr.split(' ')[0]
+            cmd = cmdutils.cmd_dict.get(cmdname)
+            return cmd and cmd.takes_count()
+
         bindings_dict = config.key_instance.get_bindings_for(modename)
         bindings = [(k, v) for (k, v) in sorted(bindings_dict.items())
                     if k.startswith(prefix) and
                     not utils.is_special_key(k) and
-                    not blacklisted(k)]
+                    not blacklisted(k) and
+                    (takes_count(v) or not countstr)]
 
         if not bindings:
             self._show_timer.stop()
diff --git a/tests/helpers/stubs.py b/tests/helpers/stubs.py
index 0bab0ce96..1d8d547fc 100644
--- a/tests/helpers/stubs.py
+++ b/tests/helpers/stubs.py
@@ -336,6 +336,7 @@ class FakeCommand:
     deprecated = attr.ib(False)
     completion = attr.ib(None)
     maxsplit = attr.ib(None)
+    takes_count = attr.ib(lambda: False)
 
 
 class FakeTimer(QObject):
diff --git a/tests/unit/misc/test_keyhints.py b/tests/unit/misc/test_keyhints.py
index 9c5c735a2..3e94f1271 100644
--- a/tests/unit/misc/test_keyhints.py
+++ b/tests/unit/misc/test_keyhints.py
@@ -92,15 +92,20 @@ def test_suggestions(keyhint, config_stub):
         ('a', 'yellow', 'c', 'message-info cmd-ac'))
 
 
-def test_suggestions_with_count(keyhint, config_stub):
-    """Test that keyhints are shown based on a prefix."""
-    bindings = {'normal': {'aa': 'message-info cmd-aa'}}
+def test_suggestions_with_count(keyhint, config_stub, monkeypatch, stubs):
+    """Test that a count prefix filters out commands that take no count."""
+    monkeypatch.setattr('qutebrowser.commands.cmdutils.cmd_dict', {
+        'foo': stubs.FakeCommand(name='foo', takes_count=lambda: False),
+        'bar': stubs.FakeCommand(name='bar', takes_count=lambda: True),
+    })
+
+    bindings = {'normal': {'aa': 'foo', 'ab': 'bar'}}
     config_stub.val.bindings.default = bindings
     config_stub.val.bindings.commands = bindings
 
     keyhint.update_keyhint('normal', '2a')
     assert keyhint.text() == expected_text(
-        ('a', 'yellow', 'a', 'message-info cmd-aa'),
+        ('a', 'yellow', 'b', 'bar'),
     )