lemonbar-driver.py

#!/usr/bin/env python

import datetime
import subprocess
import threading
import time
import os
import re

import i3ipc
import psutil
from Xlib import display, rdb, Xatom

# Get Xresources
dbstr = display.Display().screen(0).root.get_full_property(
    Xatom.RESOURCE_MANAGER, Xatom.STRING
    )
db = rdb.ResourceDB(string=dbstr.value.decode())
DPI = int(db.get("lemonbar.dpi", "lemonbar.dpi", "96"))
ALPHA = db.get("lemonbar.alpha_hex", "lemonbar.alpha_hex", "FF")
BG = db.get("lemonbar.background", "lemonbar.background", "#000000")
BG = f"#{ALPHA}{BG[1:]}"
FG = db.get("lemonbar.foreground", "lemonbar.foreground", "#FFFFFF")
COLORS = [db.get(f"lemonbar.color{i}", f"lemonbar.color{i}", "#888888")
    for i in range(16)]

BAR_HEIGHT = str(round(20*DPI/96 + 0.001))
UNDERLINE_HEIGHT = str(round(2*DPI/96 + 0.001))

BAR_COMMAND = ["lemonbar",
    "-b", "-a", "20",
    "-B", BG, "-F", FG,
    "-u", UNDERLINE_HEIGHT,
    "-f", "sans-serif:size=9",
    "-f", "Noto Sans Mono CJK JP:size=8",
    "-f", "Font Awesome 5 Free:style=Solid:size=9",
    ]

OP_COMMAND = "/usr/bin/i3-msg focus output"
WS_COMMAND = "/usr/bin/i3-msg workspace"

WS_COLOR_FOCUS = (COLORS[10], COLORS[15], COLORS[14])
WS_COLOR_NOFOC = (COLORS[6], COLORS[15], COLORS[14])
WS_COLOR_URG = (COLORS[1], COLORS[15], COLORS[9])
WS_COLOR_DEAD = ("", "", "#444444")

STAT_COLOR_NORM = (COLORS[12], None)
STAT_COLOR_HIGH = (COLORS[10], COLORS[11])
STAT_COLOR_URG = (COLORS[1], COLORS[9])

WS_TEMPL_HERE = ("%{{A:{oc} '{out}'; {wc} '{name}':}}"
    "%{{F{clr[1]}}}%{{B{clr[0]}}}\u2007{icon}\u2007%{{F-}}%{{B-}}"
    "%{{A}}")
WS_TEMPL_OFF = ("%{{A:{oc} '{out}'; {wc} '{name}':}}"
    "%{{F{clr[2]}}}%{{U{clr[0]}}}%{{+u}}\u2007{icon}\u2007%{{-u}}%{{F-}}%{{U-}}"
    "%{{A}}")
WS_TEMPL_HID = ("%{{A:{oc} '{out}'; {wc} '{name}':}}"
    "%{{F{clr[2]}}}\u2007{icon}\u2007%{{F-}}"
    "%{{A}}")

# MAILDIRS = ["~/Mail/Redacted/INBOX", "~/Mail/Redacted2/INBOX"]
MAILDIRS = []
MAILDIRS = [os.path.expanduser(p) for p in MAILDIRS]

WORKSPACES = {str(i): f"%{{T2}}{chr(0x215f + i)}%{{T-}}" for i in range(1, 11)}

SIGIL = f"%{{T2}}%{{F{COLORS[13]}}}\u273F%{{F-}}%{{T-}}"


class Output(object):

  """Class representing a display. Each Output object runs its own instance of
  lemonbar."""

  def __init__(self, data, manager):
    """Creates the Output object based on data from i3.

    :data: The dict representing the output as obtaiend from i3.
    :manager: The Manager instance that created this.

    """
    self.name = data.name
    self.bar = None
    self.workspaces = ""
    self.manager = manager

    # Start the bar process.
    r = data.rect
    geometry = f"{r.width}x{BAR_HEIGHT}+{r.x}+{r.y}"
    self.bar = subprocess.Popen(
        [*BAR_COMMAND, "-g", geometry],
        stdin=subprocess.PIPE, stdout=self.manager.shell.stdin
        )

  def write_bar(self):
    """Writes workspace and status strings to the bar."""
    if not self.bar:
      return
    output = (f"\u2007{self.workspaces}"
              f"%{{c}}{SIGIL}"
              f"%{{r}}{self.manager.status}\u2007\n")
    try:
      self.bar.stdin.write(output.encode("utf-8"))
      self.bar.stdin.flush()
    except BrokenPipeError:
      try:
        self.manager.outputs.remove(self)
      except ValueError:
        pass
      self.kill()

  def update_workspaces(self, data):
    """Updates the workspace string.

    :data: The list representing the workspaces as obtained from i3.

    """
    self.workspaces = ""

    for name, icon in WORKSPACES.items():
      self.format_workspace(name, icon, data)

    # Show any other workspaces.
    for ws in data:
      if ws.name not in WORKSPACES:
        self.format_workspace(ws.name, ws.name, data)

    self.write_bar()

  def format_workspace(self, name, icon, data):
    """Add a workspace to the status line."""
    ws = next((w for w in data if w.name == name), None)
    if not ws:
      color = WS_COLOR_DEAD
      template = WS_TEMPL_HID
    else:
      if ws.urgent:
        color = WS_COLOR_URG
      elif ws.focused:
        color = WS_COLOR_FOCUS
      else:
        color = WS_COLOR_NOFOC
      if not ws.visible:
        template = WS_TEMPL_HID
      elif ws.output == self.name:
        template = WS_TEMPL_HERE
      else:
        template = WS_TEMPL_OFF

    self.workspaces += template.format(name=name, clr=color, wc=WS_COMMAND,
        oc=OP_COMMAND, out=self.name, icon=icon)

  def kill(self):
    """Destroys the output."""
    if self.bar:
      self.bar.terminate()


class Manager(object):

  """The class responsible for managing everything."""

  def __init__(self):
    """Set up the class, start outputs, shell and subscriptions."""
    self.outputs = []
    self.conn = i3ipc.Connection()
    self.status = ""
    self.shell = subprocess.Popen(["sh", "-"], stdin=subprocess.PIPE)

    self.update_outputs()

    thr = threading.Thread(target=self.status_thread)
    thr.start()

    self.conn.on("workspace", self.update_workspaces)
    self.conn.on("output", self.update_outputs)
    self.conn.main()

  def status_thread(self):
    """Calls update_status periodically."""
    while True:
      self.update_status()
      time.sleep(5)

  def update_outputs(self, conn=None, evt=None):
    """Destroy all outputs and start new ones."""
    data = self.conn.get_outputs()

    for o in self.outputs:
      o.kill()

    self.outputs = []

    for d in data:
      if d.active:
        self.outputs.append(Output(d, self))

    self.update_workspaces()

  def update_workspaces(self, conn=None, evt=None):
    """Update the workspaces of all outputs."""
    data = self.conn.get_workspaces()
    for o in self.outputs:
      o.update_workspaces(data)

  def update_status(self):
    """Updates the status string."""
    status = []

    # mail
    # Unfortunately, the mailbox module uses too much CPU.
    unread = sum(len(os.listdir(f"{mb}/new")) for mb in MAILDIRS)
    unread += sum(not re.search('S[A-Z]*$', m)
        for mb in MAILDIRS for m in os.listdir(f"{mb}/cur"))
    if unread > 0:
      status.append(status_item(0xF0E0, f"{unread:\u2007>2d}",
          color=STAT_COLOR_HIGH))

    # wireless
    wl = next((v for k, v in psutil.net_if_addrs().items() if k[:2] == "wl"),
        None)
    if wl and len(wl) <= 1:
      status.append(status_item(0xF1EB, color=STAT_COLOR_URG))

    # battery
    bat = psutil.sensors_battery()
    if bat:
      if bat.power_plugged or bat.power_plugged is None:
        icon = 0xF376
      else:
        icon = 0xF244 - min(4, max(0, int(bat.percent/20)))
      color = STAT_COLOR_NORM
      if bat.percent <= 20:
        color = STAT_COLOR_URG
      status.append(status_item(icon, f"{bat.percent:\u2007>3.0f}%",
          color=color))

    # cpu usage
    cpu = sum(psutil.cpu_percent(percpu=True))
    color = STAT_COLOR_NORM
    if cpu >= 350:
      color = STAT_COLOR_URG
    elif cpu >= 200:
      color = STAT_COLOR_HIGH
    status.append(status_item(0xF085, f"{cpu:\u2007>3.0f}%", color=color))

    # memory usage
    mem = psutil.virtual_memory()
    used = (mem.total - mem.available)/(1<<30)
    color = STAT_COLOR_NORM
    if mem.available <= 1<<28:  # 250 MB
      color = STAT_COLOR_URG
    elif mem.available <= 1<<30:  # 1 GB
      color = STAT_COLOR_HIGH
    status.append(status_item(0xF2DB, f"{used:\u2007>4.1f}G", color=color))

    # cpu temperature
    temp = max(t.current for t in psutil.sensors_temperatures()["coretemp"])
    color = STAT_COLOR_NORM
    if temp >= 80:
      color = STAT_COLOR_URG
    elif temp >= 70:
      color = STAT_COLOR_HIGH
    icon = 0xF2CB - min(4, max(0, int((temp-30)/10)))
    status.append(status_item(icon, f"{temp:\u2007>3.0f}°C", color=color))

    # date and time
    now = datetime.datetime.now()
    status.append(status_item(0xF073, f" {now:%d %b}"))
    status.append(status_item(0xF017, f" {now:%H:%M}"))

    self.status = "".join(status)
    for o in self.outputs:
      o.write_bar()


def status_item(icon, text="", *, color=STAT_COLOR_NORM):
  """Format an item for the status line."""
  left = right = ""
  left += f"%{{U{color[0]}}}%{{+u}} "
  right = f" %{{-u}}" + right
  if color[1]:
    left += f"%{{F{color[1]}}}"
    right = f"%{{F-}}" + right
  return f"\u2007{left}{chr(icon)}{text}{right}"


if __name__ == "__main__":
  Manager()