From 8f4a21855f38650e3e3f2c5afa0aa95a3252e45c Mon Sep 17 00:00:00 2001
From: Justin Burkett <justin@burkett.cc>
Date: Wed, 30 Nov 2016 18:52:52 -0500
Subject: [PATCH] Support functions in -add-key-based-replacements

---
 which-key-tests.el | 24 +++++++++++++++++++++---
 which-key.el       | 19 ++++++++++++++-----
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/which-key-tests.el b/which-key-tests.el
index 4fcbf395385..93e1dfa969f 100644
--- a/which-key-tests.el
+++ b/which-key-tests.el
@@ -49,12 +49,24 @@
   "Test `which-key--maybe-replace'. See #154"
   (let ((which-key-replacement-alist
          '((("C-c [a-d]" . nil) . ("C-c a" . "c-c a"))
-           (("C-c .+" . nil) . ("C-c *" . "c-c *")))))
+           (("C-c .+" . nil) . ("C-c *" . "c-c *"))))
+        (test-mode-1 t)
+        (test-mode-2 nil))
     (which-key-add-key-based-replacements
       "C-c ." "test ."
       "SPC ." "SPC ."
       "C-c \\" "regexp quoting"
-      "C-c [" "bad regexp")
+      "C-c [" "bad regexp"
+      "SPC t1" (lambda (kb)
+                 (cons (car kb)
+                       (if test-mode-1
+                           "[x] test mode"
+                         "[ ] test mode")))
+      "SPC t2" (lambda (kb)
+                 (cons (car kb)
+                       (if test-mode-2
+                           "[x] test mode"
+                         "[ ] test mode"))))
     (should (equal
              (which-key--maybe-replace '("C-c g" . "test"))
              '("C-c *" . "c-c *")))
@@ -77,7 +89,13 @@
     ;; see #155
     (should (equal
              (which-key--maybe-replace '("SPC . ." . "don't replace"))
-             '("SPC . ." . "don't replace")))))
+             '("SPC . ." . "don't replace")))
+    (should (equal
+             (which-key--maybe-replace '("SPC t 1" . "test mode"))
+             '("SPC t 1" . "[x] test mode")))
+    (should (equal
+             (which-key--maybe-replace '("SPC t 2" . "test mode"))
+             '("SPC t 2" . "[ ] test mode")))))
 
 (provide 'which-key-tests)
 ;;; which-key-tests.el ends here
diff --git a/which-key.el b/which-key.el
index d383866b7d6..63327cb9ba2 100644
--- a/which-key.el
+++ b/which-key.el
@@ -796,10 +796,13 @@ may either be a string, as in
 
 \(which-key-add-key-based-replacements \"C-x 1\" \"maximize\"\)
 
-or a cons of two strings as in
+a cons of two strings as in
 
 \(which-key-add-key-based-replacements \"C-x 8\" '(\"unicode\" . \"Unicode keys\")\)
 
+or a function that takes a \(KEY . BINDING\) cons and returns a
+replacement.
+
 In the second case, the second string is used to provide a longer
 name for the keys under a prefix.
 
@@ -809,9 +812,12 @@ replacements are added to
   ;; TODO: Make interactive
   (while key-sequence
     ;; normalize key sequences before adding
-    (let ((key-seq (key-description (kbd key-sequence))))
+    (let ((key-seq (key-description (kbd key-sequence)))
+          (replace (or (and (functionp replacement) replacement)
+                       (cdr-safe replacement)
+                       replacement)))
       (push (cons (cons (concat "\\`" (regexp-quote key-seq) "\\'") nil)
-                  (cons nil (or (car-safe replacement) replacement)))
+                  (if (functionp replace) replace (cons nil replace)))
             which-key-replacement-alist)
       (when (consp replacement)
         (push (cons key-seq (cdr-safe replacement))
@@ -835,9 +841,12 @@ addition KEY-SEQUENCE REPLACEMENT pairs) to apply."
          (or (cdr-safe (assq mode which-key--prefix-title-alist)) (list))))
     (while key-sequence
     ;; normalize key sequences before adding
-      (let ((key-seq (key-description (kbd key-sequence))))
+      (let ((key-seq (key-description (kbd key-sequence)))
+            (replace (or (and (functionp replacement) replacement)
+                         (cdr-safe replacement)
+                         replacement)))
         (push (cons (cons (concat "\\`" (regexp-quote key-seq) "\\'") nil)
-                    (cons nil (or (car-safe replacement) replacement)))
+                    (if (functionp replace) replace (cons nil replace)))
               mode-alist)
         (when (consp replacement)
           (push (cons key-seq (cdr-safe replacement))