dotfiles/config/lemonbar/display.py

#!/usr/bin/env python3

import enum
import threading
import time
import subprocess
import logging
import coloredlogs

coloredlogs.install(level='DEBUG', fmt='%(levelname)s %(message)s')
log = logging.getLogger()

# TODO Allow deletion of Bar, BarGroup and Section for screen changes
# IDEA Use i3 ipc events rather than relying on xrandr or Xlib (less portable
# but easier)
# TODO Optimize to use write() calls instead of string concatenation (writing
# BarGroup strings should be a good compromise)
# TODO Use bytes rather than strings
# TODO Use default colors of lemonbar sometimes
# TODO Mouse actions
# TODO Adapt bar height with font height


class BarGroupType(enum.Enum):
    LEFT = 0
    RIGHT = 1
    # TODO Middle
    # MID_LEFT = 2
    # MID_RIGHT = 3


class BarStdoutThread(threading.Thread):
    def run(self):
        while True:
            handle = Bar.process.stdout.readline().strip()
            if not len(handle):
                continue
            if handle not in Bar.actionsH2F:
                log.error("Unknown action: {}".format(handle))
                continue
            function = Bar.actionsH2F[handle]
            function()


class Bar:
    """
    One bar for each screen
    """

    # Constants
    FONTS = ["DejaVu Sans Mono for Powerline", "Font Awesome"]
    FONTSIZE = 10

    @staticmethod
    def init():
        Section.init()

        cmd = ['lemonbar', '-b']
        for font in Bar.FONTS:
            cmd += ["-f", "{}:size={}".format(font, Bar.FONTSIZE)]
        Bar.process = subprocess.Popen(cmd, stdin=subprocess.PIPE,
                                       stdout=subprocess.PIPE)
        Bar.stdoutThread  = BarStdoutThread()
        Bar.stdoutThread.start()

        # Debug
        Bar(0)
        # Bar(1)

    # Class globals
    everyone = set()
    string = ""
    process = None

    nextHandle = 0
    actionsF2H = dict()
    actionsH2F = dict()

    @staticmethod
    def getFunctionHandle(function):
        assert callable(function)
        if function in Bar.actionsF2H.keys():
            return Bar.actionsF2H[function]

        handle = '{:x}'.format(Bar.nextHandle).encode()
        Bar.nextHandle += 1

        Bar.actionsF2H[function] = handle
        Bar.actionsH2F[handle] = function

        log.debug("Registered action {} → {}".format(handle, function))
        return handle

    @staticmethod
    def forever():
        while True:
            time.sleep(60)

    def __init__(self, screen):
        assert isinstance(screen, int)
        self.screen = "%{S" + str(screen) + "}"
        self.groups = dict()

        for groupType in BarGroupType:
            group = BarGroup(groupType, self)
            self.groups[groupType] = group

        self.childsChanged = False

        self.everyone.add(self)

    @staticmethod
    def addSectionAll(section, group, screens=None):
        """
        .. note::
            Add the section before updating it for the first time.
        """
        assert isinstance(section, Section)
        assert isinstance(group, BarGroupType)
        # TODO screens selection
        for bar in Bar.everyone:
            bar.addSection(section, group=group)

    def addSection(self, section, group):
        assert isinstance(section, Section)
        assert isinstance(group, BarGroupType)
        self.groups[group].addSection(section)

    def update(self):
        if self.childsChanged:
            self.string = self.screen
            self.string += self.groups[BarGroupType.LEFT].string
            self.string += self.groups[BarGroupType.RIGHT].string

            self.childsChanged = False

    @staticmethod
    def updateAll():

        Bar.string = ""
        for bar in Bar.everyone:
            bar.update()
            Bar.string += bar.string
        # Color for empty sections
        Bar.string += BarGroup.color(*Section.EMPTY)

        # print(Bar.string)
        Bar.process.stdin.write(bytes(Bar.string + '\n', 'utf-8'))
        Bar.process.stdin.flush()


class BarGroup:
    """
    One for each group of each bar
    """

    everyone = set()

    def __init__(self, groupType, parent):
        assert isinstance(groupType, BarGroupType)
        assert isinstance(parent, Bar)

        self.groupType = groupType
        self.parent = parent

        self.sections = list()
        self.string = ''
        self.parts = []

        #: One of the sections that had their theme or visibility changed
        self.childsThemeChanged = False

        #: One of the sections that had their text (maybe their size) changed
        self.childsTextChanged = False

        BarGroup.everyone.add(self)

    def addSection(self, section):
        self.sections.append(section)
        section.addParent(self)

    def addSectionAfter(self, sectionRef, section):
        index = self.sections.index(sectionRef)
        self.sections.insert(index + 1, section)
        section.addParent(self)

    ALIGNS = {BarGroupType.LEFT: "%{l}", BarGroupType.RIGHT: "%{r}"}

    @staticmethod
    def fgColor(color):
        return "%{F" + (color or '-') + "}"

    @staticmethod
    def bgColor(color):
        return "%{B" + (color or '-') + "}"

    @staticmethod
    def underlineColor(color):
        return "%{U" + (color or '-') + "}"

    @staticmethod
    def color(fg, bg):
        return BarGroup.fgColor(fg) + BarGroup.bgColor(bg)

    def update(self):
        if self.childsThemeChanged:
            parts = [BarGroup.ALIGNS[self.groupType]]

            secs = [sec for sec in self.sections if sec.visible]
            lenS = len(secs)
            for s in range(lenS):
                sec = secs[s]
                theme = Section.THEMES[sec.theme]
                if self.groupType == BarGroupType.LEFT:
                    oSec = secs[s + 1] if s < lenS - 1 else None
                else:
                    oSec = secs[s - 1] if s > 0 else None
                oTheme = Section.THEMES[oSec.theme] \
                    if oSec is not None else Section.EMPTY

                if self.groupType == BarGroupType.LEFT:
                    if s == 0:
                        parts.append(BarGroup.bgColor(theme[1]))
                    parts.append(BarGroup.fgColor(theme[0]))
                    parts.append(sec)
                    if theme == oTheme:
                        parts.append("")
                    else:
                        parts.append(BarGroup.color(theme[1], oTheme[1]) + "")
                else:
                    if theme is oTheme:
                        parts.append("")
                    else:
                        parts.append(BarGroup.fgColor(theme[1]) + "")
                        parts.append(BarGroup.color(*theme))
                    parts.append(sec)


            # TODO OPTI Concatenate successive strings
            self.parts = parts

        if self.childsTextChanged or self.childsThemeChanged:
            self.string = ""
            for part in self.parts:
                if isinstance(part, str):
                    self.string += part
                elif isinstance(part, Section):
                    self.string += part.curText

            self.parent.childsChanged = True

            self.childsThemeChanged = False
            self.childsTextChanged = False

    @staticmethod
    def updateAll():
        for group in BarGroup.everyone:
            group.update()
        Bar.updateAll()


class SectionThread(threading.Thread):
    ANIMATION_START = 0.025
    ANIMATION_STOP = 0.001
    ANIMATION_EVOLUTION = 0.9
    def run(self):
        while Section.somethingChanged.wait():
            Section.updateAll()
            animTime = self.ANIMATION_START
            frameTime = time.perf_counter()
            while len(Section.sizeChanging) > 0:
                frameTime += animTime
                curTime = time.perf_counter()
                sleepTime = frameTime - curTime
                time.sleep(sleepTime if sleepTime > 0 else 0)
                Section.updateAll()
                animTime *= self.ANIMATION_EVOLUTION
                if animTime < self.ANIMATION_STOP:
                    animTime = self.ANIMATION_STOP


class Section:

    # TODO Update all of that to base16
    COLORS = ['#002b36', '#dc322f', '#859900', '#b58900', '#268bd2', '#6c71c4',
              '#2aa198', '#93a1a1', '#657b83', '#dc322f', '#859900', '#b58900',
              '#268bd2', '#6c71c4', '#2aa198', '#fdf6e3']
    FGCOLOR = '#93a1a1'
    BGCOLOR = '#002b36'

    THEMES = list()
    EMPTY = (FGCOLOR, BGCOLOR)

    #: Sections that do not have their destination size
    sizeChanging = set()
    updateThread = SectionThread(daemon=True)
    somethingChanged = threading.Event()
    lastChosenTheme = 0

    @staticmethod
    def init():
        for t in range(8, 16):
            Section.THEMES.append((Section.COLORS[0], Section.COLORS[t]))

        Section.updateThread.start()

    def __init__(self, theme=None):
        #: Displayed section
        #: Note: A section can be empty and displayed!
        self.visible = False

        if theme is None:
            theme = Section.lastChosenTheme
            Section.lastChosenTheme = (Section.lastChosenTheme + 1) \
                % len(Section.THEMES)
        self.theme = theme

        #: Displayed text
        self.curText = ''
        #: Displayed text size
        self.curSize = 0

        #: Destination text
        self.dstText = ''
        #: Destination size
        self.dstSize = 0

        #: Groups that have this section
        self.parents = set()

        self.actions = set()

    def __str__(self):
        try:
            return "<{}><{}>{:01d}{}{:02d}/{:02d}" \
                .format(self.curText, self.dstText,
                        self.theme, "+" if self.visible else "-",
                        self.curSize, self.dstSize)
        except:
            return super().__str__()

    def addParent(self, parent):
        self.parents.add(parent)

    def appendAfter(self, section):
        assert len(self.parents)
        for parent in self.parents:
            parent.addSectionAfter(self, section)

    def informParentsThemeChanged(self):
        for parent in self.parents:
            parent.childsThemeChanged = True

    def informParentsTextChanged(self):
        for parent in self.parents:
            parent.childsTextChanged = True

    def parseParts(self, parts, bit='', clo=''):
        # TODO OPTI Shall benefit of a major refactor, using Objects rather
        # than dicts and a stack-based approac
        if isinstance(parts, str):
            parts = [parts]

        bits = []
        clos = []
        for part in parts:
            if isinstance(part, str):
                for char in part:
                    bit += char
                    bits.append(bit)
                    clos.append(clo)
                    bit = ''
            elif isinstance(part, dict):
                newBit = ''
                newClo = ''
                if 'fgColor' in part:
                    newBit = newBit + BarGroup.fgColor(part['fgColor'])
                    newClo = BarGroup.fgColor(Section.THEMES[self.theme][0]) + newClo
                if 'bgColor' in part:
                    newBit = newBit + BarGroup.bgColor(part['bgColor'])
                    newClo = BarGroup.bgColor(Section.THEMES[self.theme][0]) + newClo
                if 'underlineColor' in part:
                    newBit = newBit + BarGroup.underlineColor(part['underlineColor'])
                    newClo = BarGroup.underlineColor(None) + newClo
                if 'underline' in part:
                    newBit = newBit + '%{+u}'
                    newClo = '%{-u}' + newClo
                if 'overline' in part:
                    newBit = newBit + '%{+o}'
                    newClo = '%{-o}' + newClo
                if 'mouseLeft' in part:
                    handle = Bar.getFunctionHandle(part['mouseLeft'])
                    newBit = newBit + '%{A1:' + handle.decode() + ':}'
                    newClo = '%{A1}' + newClo
                if 'mouseMiddle' in part:
                    handle = Bar.getFunctionHandle(part['mouseMiddle'])
                    newBit = newBit + '%{A2:' + handle.decode() + ':}'
                    newClo = '%{A2}' + newClo
                if 'mouseRight' in part:
                    handle = Bar.getFunctionHandle(part['mouseRight'])
                    newBit = newBit + '%{A3:' + handle.decode() + ':}'
                    newClo = '%{A3}' + newClo
                if 'mouseScrollUp' in part:
                    handle = Bar.getFunctionHandle(part['mouseScrollUp'])
                    newBit = newBit + '%{A4:' + handle.decode() + ':}'
                    newClo = '%{A4}' + newClo
                if 'mouseScrollDown' in part:
                    handle = Bar.getFunctionHandle(part['mouseScrollDown'])
                    newBit = newBit + '%{A5:' + handle.decode() + ':}'
                    newClo = '%{A5}' + newClo
                newBits, newClos = self.parseParts(part["cont"], bit=bit+newBit, clo=newClo+clo)
                newBits[-1] += newClo # TODO Will sometimes display errors from lemonbar
                bits += newBits
                clos += newClos
            else:
                raise RuntimeError()

        return bits, clos

    def updateText(self, text):
        # TODO FEAT Actions
        # TODO OPTI When srcSize == dstSize, maybe the bit array isn't
        # needed
        oldActions = self.actions.copy()

        if len(text):
            if isinstance(text, str):
                # TODO OPTI This common case
                text = [text]

            self.dstBits, self.dstClos = self.parseParts([' '] + text + [' '])
            # TODO FEAT Half-spaces

            self.dstText = ''.join(self.dstBits)
            self.dstSize = len(self.dstBits)
        else:
            self.dstSize = 0

        for action in oldActions:
            self.unregsiterAction(action)

        if self.curSize == self.dstSize:
            self.curText = self.dstText
            self.informParentsTextChanged()
        else:
            Section.sizeChanging.add(self)
        Section.somethingChanged.set()

    def updateTheme(self, theme):
        assert isinstance(theme, int)
        assert theme < len(Section.THEMES)
        self.theme = theme
        self.informParentsThemeChanged()
        Section.somethingChanged.set()

    def updateVisibility(self, visibility):
        assert isinstance(visibility, bool)

        self.visible = visibility
        self.informParentsThemeChanged()
        Section.somethingChanged.set()

    @staticmethod
    def fit(text, size):
        t = len(text)
        return text[:size] if t >= size else text + [" "] * (size - t)

    def update(self):
        # TODO Might profit of a better logic
        if not self.visible:
            self.updateVisibility(True)
            return

        if self.dstSize > self.curSize:
            self.curSize += 1
        elif self.dstSize < self.curSize:
            self.curSize -= 1
        else:
            # Visibility toggling must be done one step after curSize = 0
            if self.dstSize == 0:
                self.updateVisibility(False)
            Section.sizeChanging.remove(self)
            return

        closPos = self.curSize-1
        clos = self.dstClos[closPos] if closPos < len(self.dstClos) else ''
        self.curText = ''.join(Section.fit(self.dstBits, self.curSize)) + clos
        self.informParentsTextChanged()

    @staticmethod
    def updateAll():
        """
        Process all sections for text size changes
        """

        for sizeChanging in Section.sizeChanging.copy():
            sizeChanging.update()

        BarGroup.updateAll()

        Section.somethingChanged.clear()

    @staticmethod
    def ramp(p, ramp=" ▁▂▃▄▅▆▇█"):
        if p > 1:
            return ramp[-1]
        elif p < 0:
            return ramp[0]
        else:
            return ramp[round(p * (len(ramp)-1))]


class StatefulSection(Section):
    # TODO Allow to temporary expand the section (e.g. when important change)
    NUMBER_STATES = 4

    def __init__(self, *args, **kwargs):
        Section.__init__(self, *args, **kwargs)
        self.state = 0
        if hasattr(self, 'onChangeState'):
            self.onChangeState(self.state)

    def incrementState(self):
        newState = min(self.state + 1, self.NUMBER_STATES - 1)
        self.changeState(newState)

    def decrementState(self):
        newState = max(self.state - 1, 0)
        self.changeState(newState)

    def changeState(self, state):
        assert isinstance(state, int)
        assert state < self.NUMBER_STATES
        self.state = state
        if hasattr(self, 'onChangeState'):
            self.onChangeState(state)
        self.refreshData()

    def updateText(self, text):
        if not len(text):
            return text
        Section.updateText(self, [{
            "mouseScrollUp": self.incrementState,
            "mouseScrollDown": self.decrementState,
            "cont": text
        }])

class ColorCountsSection(StatefulSection):
    NUMBER_STATES = 3
    ICON = '?'

    def __init__(self, theme=None):
        StatefulSection.__init__(self, theme=theme)

    def fetcher(self):
        counts = self.subfetcher()
        # Nothing
        if not len(counts):
            return ''
        # Icon colored
        elif self.state == 0 and len(counts) == 1:
            count, color = counts[0]
            return [{'fgColor': color, "cont": self.ICON}]
        # Icon
        elif self.state == 0 and len(counts) > 1:
            return self.ICON
        # Icon + Total
        elif self.state == 1 and len(counts) > 1:
            total = sum([count for count, color in counts])
            return [self.ICON, ' ', str(total)]
        # Icon + Counts
        else:
            text = [self.ICON]
            for count, color in counts:
                text += [' ', {'fgColor': color, "cont": str(count)}]
            return text


if __name__ == '__main__':
    Bar.init()
    sec = Section(0)
    sech = Section(1)
    sec1 = Section(2)
    sec2 = Section(3)
    sec2h = Section(4)
    sec3 = Section(5)
    Bar.addSectionAll(sec, BarGroupType.LEFT)
    Bar.addSectionAll(sech, BarGroupType.LEFT)
    Bar.addSectionAll(sec1, BarGroupType.LEFT)
    Bar.addSectionAll(sec2, BarGroupType.RIGHT)
    Bar.addSectionAll(sec2h, BarGroupType.RIGHT)
    Bar.addSectionAll(sec3, BarGroupType.RIGHT)

    time.sleep(1)
    sec.updateText("A")
    time.sleep(1)
    sec.updateText("")
    time.sleep(1)
    sec.updateText("Hello")
    sec1.updateText("world!")
    sec2.updateText("Salut")
    sec2h.updateText("le")
    sec3.updateText("monde !")
    time.sleep(3)
    sech.updateText("the")
    sec2h.updateText("")
    time.sleep(2)
    sec.updateText("")
    sech.updateText("")
    sec1.updateText("")
    sec2.updateText("")
    sec2h.updateText("")
    sec3.updateText("")
    time.sleep(5)