Add support for XEP-0392 (Consistent Color Generation)

4 files changed, 151 insertions, 27 deletions

diff --git a/poezio/colors.py b/poezio/colors.py
new file mode 100644
index 00000000..197120ad
--- /dev/null
+++ b/poezio/colors.py
@@ -0,0 +1,102 @@
+import curses
+import hashlib
+import math
+
+# BT.601 (YCbCr) constants, see XEP-0392
+K_R = 0.299
+K_G = 0.587
+K_B = 1-K_R-K_G
+
+def ncurses_color_to_rgb(color):
+    if color <= 15:
+        try:
+            (r, g, b) = curses.color_content(color)
+        except:  # fallback in faulty terminals (e.g. xterm)
+            (r, g, b) = curses.color_content(color%8)
+        r = r / 1000 * 6 - 0.01
+        g = g / 1000 * 6 - 0.01
+        b = b / 1000 * 6 - 0.01
+    elif color <= 231:
+        color = color - 16
+        r = color % 6
+        color = color / 6
+        g = color % 6
+        color = color / 6
+        b = color % 6
+    else:
+        color -= 232
+        r = g = b = color / 24 * 6
+    return r / 6, g / 6, b / 6
+
+def rgb_to_ycbcr(r, g, b):
+    y = K_R * r + K_G * g + K_B * b
+    cr = (r - y) / (1 - K_R) / 2
+    cb = (b - y) / (1 - K_B) / 2
+    return y, cb, cr
+
+def generate_ccg_palette(curses_palette, reference_y):
+    cbcr_palette = {}
+    for curses_color in curses_palette:
+        r, g, b = ncurses_color_to_rgb(curses_color)
+        # drop grayscale
+        if r == g == b:
+            continue
+        y, cb, cr = rgb_to_ycbcr(r, g, b)
+        key = round(cbcr_to_angle(cb, cr), 2)
+        try:
+            existing_y, *_ = cbcr_palette[key]
+        except KeyError:
+            pass
+        else:
+            if abs(existing_y - reference_y) <= abs(y - reference_y):
+                continue
+        cbcr_palette[key] = y, curses_color
+    return {
+        angle: curses_color
+        for angle, (_, curses_color) in cbcr_palette.items()
+    }
+
+def text_to_angle(text):
+    hf = hashlib.sha1()
+    hf.update(text.encode("utf-8"))
+    hue = int.from_bytes(hf.digest()[:2], "little")
+    return hue / 65535 * math.pi * 2
+
+def angle_to_cbcr_edge(angle):
+    cr = math.sin(angle)
+    cb = math.cos(angle)
+    if abs(cr) > abs(cb):
+        factor = 0.5 / abs(cr)
+    else:
+        factor = 0.5 / abs(cb)
+    return cb*factor, cr*factor
+
+def cbcr_to_angle(cb, cr):
+    magn = math.sqrt(cb**2 + cr**2)
+    if magn > 0:
+        cr /= magn
+        cb /= magn
+    return math.atan2(cr, cb) % (2*math.pi)
+
+def ccg_palette_lookup(palette, angle):
+    # try quick lookup first
+    try:
+        color = palette[round(angle, 2)]
+    except KeyError:
+        pass
+    else:
+        return color
+
+    best_metric = float("inf")
+    best = None
+    for anglep, color in palette.items():
+        metric = abs(anglep - angle)
+        if metric < best_metric:
+            best_metric = metric
+            best = color
+
+    return best
+
+def ccg_text_to_color(palette, text):
+    angle = text_to_angle(text)
+    return ccg_palette_lookup(palette, angle)
diff --git a/poezio/theming.py b/poezio/theming.py
index 49f838d2..4c93d396 100755
--- a/poezio/theming.py
+++ b/poezio/theming.py
@@ -75,6 +75,7 @@ import curses
 import functools
 import os
 from os import path
+from poezio import colors
 
 from importlib import machinery
 finder = machinery.PathFinder()
@@ -363,6 +364,10 @@ class Theme(object):
         'default': (7, -1),
     }
 
+    @property
+    def ccg_palette(self):
+        prepare_ccolor_palette(self)
+        return self.CCG_PALETTE
 
 # This is the default theme object, used if no theme is defined in the conf
 theme = Theme()
@@ -508,6 +513,30 @@ def update_themes_dir(option=None, value=None):
 
     log.debug('Theme load path: %s', load_path)
 
+def prepare_ccolor_palette(theme):
+    """
+    Prepare the Consistent Color Generation (XEP-0392) palette for a theme.
+    """
+    if hasattr(theme, "CCG_PALETTE"):
+        # theme overrides the palette
+        return
+
+    if any(bg != -1 for fg, bg in theme.LIST_COLOR_NICKNAMES):
+        # explicitly disable CCG, can’t handle dynamic background colors
+        theme.CCG_PALETTE = None
+        return
+
+    if not hasattr(theme, "CCG_Y"):
+        theme.CCG_Y = 0.5**0.45
+
+    theme.CCG_PALETTE = colors.generate_ccg_palette(
+        [
+            fg for fg, _ in theme.LIST_COLOR_NICKNAMES
+            # exclude grayscale
+            if fg < 232
+        ],
+        theme.CCG_Y,
+    )
 
 def reload_theme():
     theme_name = config.get('theme')
@@ -531,6 +560,7 @@ def reload_theme():
 
     if hasattr(new_theme, 'theme'):
         theme = new_theme.theme
+        prepare_ccolor_palette(theme)
     else:
         return 'No theme present in the theme file'
 
diff --git a/poezio/user.py b/poezio/user.py
index 68d52493..40d042f8 100644
--- a/poezio/user.py
+++ b/poezio/user.py
@@ -12,7 +12,7 @@ An user is a MUC participant, not a roster contact (see contact.py)
 from random import choice
 from datetime import timedelta, datetime
 from hashlib import md5
-from poezio import xhtml
+from poezio import xhtml, colors
 
 from poezio.theming import get_theme
 
@@ -41,6 +41,8 @@ class User(object):
         self.last_talked = datetime(1, 1, 1)  # The oldest possible time
         self.update(affiliation, show, status, role)
         self.change_nick(nick)
+        self.jid = jid
+        self.chatstate = None
         if color != '':
             self.change_color(color, deterministic)
         else:
@@ -48,14 +50,20 @@ class User(object):
                 self.set_deterministic_color()
             else:
                 self.color = choice(get_theme().LIST_COLOR_NICKNAMES)
-        self.jid = jid
-        self.chatstate = None
 
     def set_deterministic_color(self):
         theme = get_theme()
-        mod = len(theme.LIST_COLOR_NICKNAMES)
-        nick_pos = int(md5(self.nick.encode('utf-8')).hexdigest(), 16) % mod
-        self.color = theme.LIST_COLOR_NICKNAMES[nick_pos]
+        if theme.ccg_palette:
+            # use XEP-0392 CCG
+            fg_color = colors.ccg_text_to_color(
+                theme.ccg_palette,
+                self.jid.bare if self.jid and self.jid.bare else self.nick
+            )
+            self.color = fg_color, -1
+        else:
+            mod = len(theme.LIST_COLOR_NICKNAMES)
+            nick_pos = int(md5(self.nick.encode('utf-8')).hexdigest(), 16) % mod
+            self.color = theme.LIST_COLOR_NICKNAMES[nick_pos]
 
     def update(self, affiliation, show, status, role):
         self.affiliation = affiliation
diff --git a/poezio/xhtml.py b/poezio/xhtml.py
index ee98f23d..4f469bd7 100644
--- a/poezio/xhtml.py
+++ b/poezio/xhtml.py
@@ -24,6 +24,7 @@ from xml.sax import saxutils
 
 from slixmpp.xmlstream import ET
 from poezio.config import config
+from poezio.colors import ncurses_color_to_rgb
 
 digits = '0123456789'  # never trust the modules
 
@@ -212,6 +213,9 @@ def get_body_from_message_stanza(message,
     content = content if content else message['body']
     return content or " "
 
+def rgb_to_html(rgb):
+    r, g, b = rgb
+    return '#%02X%02X%02X' % (int(r*256), int(g*256), int(b*256))
 
 def ncurses_color_to_html(color):
     """
@@ -219,27 +223,7 @@ def ncurses_color_to_html(color):
     a string of the form #XXXXXX representing an
     html color.
     """
-    if color <= 15:
-        try:
-            (r, g, b) = curses.color_content(color)
-        except:  # fallback in faulty terminals (e.g. xterm)
-            (r, g, b) = curses.color_content(color % 8)
-        r = r / 1000 * 6 - 0.01
-        g = g / 1000 * 6 - 0.01
-        b = b / 1000 * 6 - 0.01
-    elif color <= 231:
-        color = color - 16
-        r = color % 6
-        color = color / 6
-        g = color % 6
-        color = color / 6
-        b = color % 6
-    else:
-        color -= 232
-        r = g = b = color / 24 * 6
-    return '#%02X%02X%02X' % (int(r * 256 / 6), int(g * 256 / 6),
-                              int(b * 256 / 6))
-
+    return rgb_to_html(ncurses_color_to_rgb(color))
 
 def _parse_css_color(name):
     if name[0] == '#':