ws_tool

#!/usr/bin/env python3
import os
import sys
import argparse
from Xlib import X
from Xlib.display import Display
from Xlib.protocol.event import ClientMessage
from Xlib.error import DisplayNameError

display_address = None
display = None
root = None
number_of_desktops = None
current_desktop = None


class Atom:
  def __init__ (self, name):
    self._name = name
    self._atom = display.intern_atom (name)

  def unwrap (self):
    return self._atom

  def get_property (self):
    p = root.get_full_property (self._atom, X.AnyPropertyType)
    if p is None:
      print (f"Property '{self._name}' not found; is this the correct display: '{display_address}'?")
      exit (1)
    return p.value[0]


class Action:
  def get ():
    print ("Number of workspaces:", number_of_desktops.get_property ())
    print ("Current workspace:   ", current_desktop.get_property () + 1)


  def set (ws):
    def do_set (idx):
      data = idx.to_bytes (4, sys.byteorder).ljust (20)
      event = ClientMessage (
        window = root,
        client_type = current_desktop.unwrap (),
        data = (32, data)
      )
      root.send_event (event, event_mask=X.SubstructureRedirectMask)
      # Without this the window manager does not receive a ClientMessage event???
      current_desktop.get_property ()
    N = number_of_desktops.get_property ()
    if ws == "next":
      idx = (current_desktop.get_property () + 1) % N
      do_set (idx)
      ws = idx + 1
    elif ws == "prev":
      idx = (current_desktop.get_property () - 1) % N
      do_set (idx)
      ws = idx + 1
    else:
      if ws - 1 not in range (N):
        print ("Workspace number must be between 1 and", N)
        return
      do_set (ws - 1)
    print (f"Switched to workspace {ws}")


  def format (inactive, active, join):
    N = number_of_desktops.get_property ()
    C = current_desktop.get_property ()
    print ((join or ' ').join (
      [inactive, active][i==C].format (w=i+1) for i in range (N)))


def parse_args ():
  p = argparse.ArgumentParser ()
  s = p.add_subparsers (dest="sub_command", title="subcommands")
  s.required = True
  # get
  get_p = s.add_parser ("get")
  # set
  set_p = s.add_parser ("set",
      epilog="If 'next' or 'prev', select the next/previous workspace.\n" +\
             "Otherwise select the workspace with the given index")
  set_p.add_argument ("workspace",
      help="Either 'next', 'prev', or a workspace index")
  # format
  fmt_p = s.add_parser ("format",
      epilog="In the format strings `{w}` is replaced with the workspace " +\
             "number and is formatted like a python format string.")
  fmt_p.add_argument ("inactive",
      help="Format string for inactive workspaces")
  fmt_p.add_argument ("active",
      help="Format string for the active workspace")
  fmt_p.add_argument ("join", nargs='?', default=' ',
      help="Separator between workspaces")
  # Process
  args = p.parse_args ()
  if args.sub_command == "set":
    if args.workspace not in ("next", "prev"):
      try:
        args.workspace = int (args.workspace)
      except ValueError:
        print ("workspace must 'next', 'prev', or a workspace index", file=sys.stderr)
        exit (1)
  return args


def main ():
  global display_address, display, root, number_of_desktops, current_desktop
  args = parse_args ()
  display_address = os.getenv ("DISPLAY")
  if display_address is None:
    print ("No display set")
    return
  try:
    display = Display (display_address)
  except DisplayNameError:
    print (f"ws_tool: Can't open display '{display_address}'")
    exit (1)
  root = display.screen ().root
  number_of_desktops = Atom ("_NET_NUMBER_OF_DESKTOPS")
  current_desktop = Atom ("_NET_CURRENT_DESKTOP")
  if args.sub_command == "get":
    Action.get ()
  elif args.sub_command == "set":
    Action.set (args.workspace)
  elif args.sub_command == "format":
    Action.format (args.inactive, args.active, args.join)

if __name__ == "__main__":
  main ()