add feature to allow manual coloring

4 files changed, 138 insertions, 10 deletions

diff --git a/README.rst b/README.rst
index 80b1387..28a192b 100644
--- a/README.rst
+++ b/README.rst
@@ -23,6 +23,13 @@ names. This means with a proper ``status-interval`` set in your tmux (see
 below), it can even automatically change when various programs in your shell
 runs!
 
+Features
+========
+
+- Support Powerline symbols
+- Automically color the window tabs by their name hash
+- Allow manual coloring
+
 Installation
 ============
 
@@ -44,7 +51,15 @@ Installation
 Customization
 =============
 
-- If you would like to use Powerline Symbols like shown in the demo, add the
+- To manually set the color of the active window tag, press ``prefix`` + C and:
+
+  - ``color-idx <256 color code>`` to manually set the color for the window index
+  - ``color-name <256 color code>`` to manually set the color for the name
+  = ``clear-idx`` clears the preivous color of the index
+  = ``clear-name`` clears the preivous color of the name
+  = ``clear-all`` use auto-coloring for all window tags
+
+- If you would like to use Powerline symbols like shown in the demo, add the
   following line to the top of your ``.tmux.conf`` to enable them:
 
   ::
diff --git a/name2color.py b/name2color.py
index ad25518..e1be6cb 100755
--- a/name2color.py
+++ b/name2color.py
@@ -1,17 +1,103 @@
 #!/usr/bin/env python
+import argparse
+import os
 import sys
 import hashlib
+import pickle
+
+
+def warn(msg):
+    sys.stderr.write(msg + '\n')
+
+
+def error(msg):
+    print(msg + '\n')
+
+
+saved_state = os.path.expanduser("~/.tmux-colortag.state")
 # add your favorite color code (256) here as candidates
 colors = [
     63, 64, 65, 68, 103, 107,
     166, 168, 172, 208, 203,
     160, 161, 167, 137, 131
 ]
-h = hashlib.sha1(sys.argv[1].encode('utf-8')).digest()
-if sys.version_info.major == 3:
-    hn = int.from_bytes(h, byteorder='big')
-else:
-    hn = 0
-    for b in bytearray(h):
-        hn = hn * 256 + int(b)
-print("colour{}".format(colors[hn % len(colors)]))
+
+parser = argparse.ArgumentParser(description='Maps tmux window tags to colors.')
+parser.add_argument('idx', type=int, help='index of the window')
+parser.add_argument('name', type=str, help='name of the window')
+parser.add_argument('--color-idx', type=int, help='manually change the color mapping for an index')
+parser.add_argument('--color-name', type=int, help='manually change the color mapping for a name')
+parser.add_argument('--clear-idx', action='store_true')
+parser.add_argument('--clear-name', action='store_true')
+parser.add_argument('--clear', action='store_true')
+args = parser.parse_args()
+
+state = {}
+changed = True
+try:
+    with open(saved_state, "rb") as f:
+        _state = pickle.load(f)
+        if type(_state) is dict:
+            state = _state
+            changed = False
+except Exception:
+    pass
+
+if args.clear:
+    state = {}
+    changed = True
+
+if len(args.name) > 100:
+    error("invalid name")
+if args.idx < 0 or args.idx >= 1024:
+    error("invalid idx")
+
+if args.clear_idx:
+    try:
+        del state[args.idx]
+        changed = True
+    except Exception:
+        pass
+
+if args.clear_name:
+    try:
+        del state[args.name]
+        changed = True
+    except Exception:
+        pass
+
+if not (args.color_idx is None):
+    if args.color_idx < 0 and args.color_idx >= 256:
+        error("invalid color code")
+    state[args.idx] = args.color_idx
+    changed = True
+
+if not (args.color_name is None):
+    if args.color_name < 0 and args.color_name >= 256:
+        error("invalid color code")
+    state[args.name] = args.color_name
+    changed = True
+
+if changed:
+    warn("wrote to the state")
+    try:
+        with open(saved_state, "wb") as f:
+            pickle.dump(state, f)
+    except Exception:
+        warn("failed to dump file")
+    sys.exit(0)
+
+c = state.get(args.idx, None)
+if c is None:
+    c = state.get(args.name, None)
+if c is None:
+    h = hashlib.sha1(args.name.encode('utf-8')).digest()
+    if sys.version_info.major == 3:
+        hn = int.from_bytes(h, byteorder='big')
+    else:
+        hn = 0
+        for b in bytearray(h):
+            hn = hn * 256 + int(b)
+    c = colors[hn % len(colors)]
+
+print("colour{}".format(c))
diff --git a/tmux-colortag-prompt.sh b/tmux-colortag-prompt.sh
new file mode 100755
index 0000000..74cfdf8
--- /dev/null
+++ b/tmux-colortag-prompt.sh
@@ -0,0 +1,24 @@
+#!/usr/bin/env bash
+
+CURRENT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+case "$1" in
+    prompt)
+        tmux command-prompt -p '[ColorTag]:' "run-shell 'idx=#I name=#W $CURRENT_DIR/tmux-colortag-prompt.sh %1'"
+        ;;
+    color-idx)
+        "$CURRENT_DIR/name2color.py" "$idx" "$name" --color-idx "$2" || echo "invalid argument"
+        ;;
+    color-name)
+        "$CURRENT_DIR/name2color.py" "$idx" "$name" --color-name "$2" || echo "invalid argument"
+        ;;
+    clear-idx)
+        "$CURRENT_DIR/name2color.py" "$idx" "$name" --clear-idx
+        ;;
+    clear-name)
+        "$CURRENT_DIR/name2color.py" "$idx" "$name" --clear-name
+        ;;
+    clear-all)
+        "$CURRENT_DIR/name2color.py" "$idx" "$name" --clear
+        ;;
+    *) echo "invalid ColorTag command"; exit 0;;
+esac
diff --git a/tmux-colortag.tmux b/tmux-colortag.tmux
index 7861bbc..4fd47b4 100755
--- a/tmux-colortag.tmux
+++ b/tmux-colortag.tmux
@@ -4,6 +4,7 @@ CURRENT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
 TMUX_COLORTAG_SET_INTERVAL="${TMUX_COLORTAG_SET_INTERVAL:-yes}"
 TMUX_COLORTAG_TAG_ONLY="${TMUX_COLORTAG_TAG_ONLY:-no}"
 TMUX_COLORTAG_USE_POWERLINE="${TMUX_COLORTAG_USE_POWERLINE:-no}"
+TMUX_COLORTAG_KEY="${TMUX_COLORTAG_KEY:-C}"
 
 if [[ "$TMUX_COLORTAG_SET_INTERVAL" == yes ]]; then
     tmux set -g status on
@@ -43,7 +44,7 @@ RIGHTBAR_DEFAULT0="#[fg=$color1,bg=$color0]"
 RIGHTBAR_HOST="#[fg=$color0,bg=$color4]"
 RIGHTBAR_HOST0="#[fg=$color4,bg=$color1]"
 LOAD_DISP="#(awk '{print \$1, \$2, \$3}' /proc/loadavg)"
-TAB_COLOR="#(\"$CURRENT_DIR/name2color.py\" #W)"
+TAB_COLOR="#(\"$CURRENT_DIR/name2color.py\" #I #W)"
 TAB_NORMAL_BEGIN="#[fg=$color0,bg=$TAB_COLOR]"
 TAB_END="#[fg=$TAB_COLOR,bg=$color0]"
 TAB_FOCUS_BEGIN_BG="#[bg=$TAB_COLOR]"
@@ -80,3 +81,5 @@ else
         "${TAB_FOCUS_BEGIN_BG}$TMUX_ARROW_SYMBOL_R1 " \
         "${TAB_FOCUS_BEGIN_FG}#I${TMUX_COLORTAG_IDX_SEP}#W${TAB_END}$TMUX_ARROW_SYMBOL_R1 ")"
 fi
+
+tmux bind-key "$TMUX_COLORTAG_KEY" run-shell "'$CURRENT_DIR/tmux-colortag-prompt.sh' prompt"