Script to change monitor config when using 4 specific monitors in a .:. configuration in x11 (terrible code)

monitor.py

"""Badly coded script to change monitor config when using 4 specific monitors in a .:. configuration in x11"""

from __future__ import annotations

import functools
import os
import re
import sys
import subprocess
import time
import warnings
from dataclasses import dataclass
from typing import Optional, TypeAlias, NamedTuple

try:
    from rich import print
except ImportError:
    pass

try:
    from rich.markup import escape
except ImportError:
    def escape(s: str) -> str:
        return s

class MonitorNames:
    LG_GL850: str = 'LG Electronics 27GL850'
    SAMSUNG_C24FG7X: str = 'Samsung Electric Company C24FG7x'
    ACER_G247HL: str = 'Acer Technologies G247HL'
    ACER_V243HL: str = 'Acer Technologies V243HL'

PHYSICAL_LOCS: dict[str, str] = {
    'LEFT': MonitorNames.ACER_G247HL,
    'TOP': MonitorNames.ACER_V243HL,
    'MID': MonitorNames.LG_GL850,
    'RIGHT': MonitorNames.SAMSUNG_C24FG7X,
}
DEFAULT_CONFIG: str = 'all'

MonitorLocations: TypeAlias = tuple[
                  tuple[None, Optional[str], None],
                  tuple[Optional[str], Optional[str], Optional[str]]
              ]

CONFIGS: dict[str, MonitorLocations] = {
    'off': (
        (None, None, None),
        (None, None, None)
    ),
    'all': (
        (None, 'TOP', None),
        ('LEFT', 'MID', 'RIGHT')
    ),
    '3x1': (
        (None, None, None),
        ('LEFT', 'MID', 'RIGHT')
    ),
    'top': (
        (None, 'TOP', None),
        (None, None, None)
    ),
    '2x1r': (
        (None, None, None),
        (None, 'MID', 'RIGHT')
    ),
    '2x1l': (
        (None, None, None),
        ('LEFT', 'MID', None)
    ),
    '2x2r': (
        (None, 'TOP', None),
        (None, 'MID', 'RIGHT')
    ),

    '2x2l': (
        (None, 'TOP', None),
        ('LEFT', 'MID', None)
    ),
    '1': (
        (None, None, None),
        (None, 'MID', None)
    ),
    '1l': (
        (None, None, None),
        ('LEFT', None, None)
    ),
    '1r': (
        (None, None, None),
        (None, None, 'RIGHT')
    ),
    '1x2': (
        (None, 'TOP', None),
        (None, 'MID', None)
    ),
}
CONFIGS['3x2'] = CONFIGS['all']
CONFIGS['0x2'] = CONFIGS['top']
CONFIGS['mid'] = CONFIGS['1']
CONFIGS['left'] = CONFIGS['1l']
CONFIGS['right'] = CONFIGS['1r']
CONFIGS['m'] = CONFIGS['1']
CONFIGS['l'] = CONFIGS['1l']
CONFIGS['r'] = CONFIGS['1r']
CONFIGS['t'] = CONFIGS['top']
CONFIGS['2r'] = CONFIGS['2x1r']
CONFIGS['2l'] = CONFIGS['2x1l']
CONFIGS['2x1'] = CONFIGS['2x1r']
CONFIGS['2x2'] = CONFIGS['2x2r']
CONFIGS['on'] = CONFIGS['all']

CONFIGS['ml'] = CONFIGS['lm'] = CONFIGS['2l']
CONFIGS['mr'] = CONFIGS['rm'] = CONFIGS['2r']
CONFIGS['tm'] = CONFIGS['mt'] = CONFIGS['1x2']

CONFIGS['lmr'] = CONFIGS['lrm'] = CONFIGS['mlr'] = CONFIGS['mrl'] = CONFIGS['rlm'] = CONFIGS['rml'] = CONFIGS['3x1']

CONFIGS['t'] = CONFIGS['top']

if DEFAULT_CONFIG not in CONFIGS:
    raise RuntimeError(f"DEFAULT_CONFIG {DEFAULT_CONFIG!r} was not found in CONFIGS...")

for val in CONFIGS.values():
    if not val[0][0] is None:
        raise RuntimeError("Upper-left monitor not present")
    if not val[0][2] is None:
        raise RuntimeError("Upper-right monitor not present")

REX_DFP_ONLY: re.Pattern = re.compile(r'DFP-\d{1}')
REX_DFP_END: re.Pattern = re.compile(r'(?P<name>.+) \((?P<number>DFP-\d{1})\)')
REX_XRANDR: re.Pattern = re.compile(
    (
        r'.*(?P<xrandrname>DP-\d{1,2}|HDMI-\d{1,2}|USB-C-\d{1,2}|VGA\d{1,2}) connected.*\s*'
        r'Identifier: (?P<id>.+)(?:[\s\S]*?)'
        r'EDID:(?:[\s\S]*?)\s+(?P<edid>[\s|\S]*?)'
        r'(?:CTM|BorderDimensions)(?:[\s\S]*?)'
        r'ConnectorType: (?P<connector>.*)\s+'
        r'ConnectorNumber: (?P<conn_num>\d{1,2})(?:[\s\S]*?)'
        r'(?P<resx>\d{1,5})x(?P<res>\d{1,5}[\s|\S]*?)'
    )
)

class Resolution(str):
    pass

class RefreshRate(str):
    pass

class Resolutions:
    VGA: Resolution = "640x480"
    SVGA: Resolution = "800x600"
    XGA: Resolution = "1024x768"
    HD720P: Resolution = "1280x720"
    LAPTOP_1366 = "1366X768"
    WSXGA_PLUS_1680: Resolution = "1680x1050"
    FULLHD: Resolution = "1920x1080"
    WQHD: Resolution = "2560x1440"
    UHD: Resolution = "3840x2160"

class RefreshRates:
    HZ60: RefreshRate = "60"
    HZ144: RefreshRate = "144"
    HZ120: RefreshRate = "120"

class EDID(str):
    pass


EDIDS: dict[str, tuple[RefreshRate, EDID]] = {
    MonitorNames.LG_GL850: (
        RefreshRate('144.00'),
        EDID(
            "00ffffffffffff001e6d7f5b37b50000"
            "061e0104b53c22789f8cb5af4f43ab26"
            "0e5054254b007140818081c0a9c0b300"
            "d1c08100d1cf28de0050a0a038500830"
            "080455502100001a000000fd003090e6"
            "e63c010a202020202020000000fc0032"
            "37474c3835300a2020202020000000ff"
            "003030364e54475931433339310a0128"
            "02031a7123090607e305c000e6060501"
            "60592846100403011f13565e00a0a0a0"
            "29503020350055502100001a909b0050"
            "a0a046500820880c555021000000b8bc"
            "0050a0a055500838f80c55502100001a"
            "00000000000000000000000000000000"
            "00000000000000000000000000000000"
            "0000000000000000000000000000001a"
        )
    ),
    MonitorNames.SAMSUNG_C24FG7X: (
        RefreshRate('144.00'),
        EDID(
            "00ffffffffffff004c2d430e50475230"
            "2f1c0104a5351e783b1c05b04e47a826"
            "105054bfef8081c0810081809500a9c0"
            "b300714f01015a8780a070384d403020"
            "350014302100001a000000fd005a90a8"
            "a823010a202020202020000000fc0043"
            "3234464737780a2020202020000000ff"
            "004854484b4230323935360a20200131"
            "020315f148903f401f04130312230907"
            "07830100000474801871382d40582c45"
            "0014302100001e047480d072382d4010"
            "2c458014302100001e023a801871382d"
            "40582c450014302100001e0000000000"
            "00000000000000000000000000000000"
            "00000000000000000000000000000000"
            "0000000000000000000000000000008a"
        )
    ),
    MonitorNames.ACER_G247HL: (
        RefreshRate('60.00'),
        EDID(
            "00ffffffffffff000472fa0337210143"
            "1e18010380351e78ca3f60a455529e27"
            "0d5054bfef80714f8140818081c08100"
            "9500b300d1c0023a801871382d40582c"
            "4500132b2100001e000000fd00384c1f"
            "5011000a202020202020000000fc0047"
            "323437484c0a202020202020000000ff"
            "005430564545303031383534300a01f4"
            "020324f14f9002030405060701111213"
            "1415161f230907078301000067030c00"
            "1000382d023a801871382d40582c4500"
            "132b2100001f011d8018711c1620582c"
            "2500132b2100009f011d007251d01e20"
            "6e285500132b2100001e8c0ad08a20e0"
            "2d10103e9600132b2100001800000000"
            "0000000000000000000000000000007e"
        )
    ),
    MonitorNames.ACER_V243HL: (
        RefreshRate('60.00'),
        EDID(
            "00ffffffffffff000472be00438e6004"
            "2e14010308351d78ea6085a6564a9c25"
            "125054b30c00714f81009500d1c08180"
            "010101010101023a801871382d40582c"
            "4500132a2100001a000000fc00563234"
            "33484c0a202020202020000000fd0037"
            "4b1e5012000a202020202020000000ff"
            "004c474b3043303136343032420a00c6"
        )
    ),
}


@functools.lru_cache(maxsize=8)
def name_from_edid(edid: EDID | str) -> str:
    for k, v in EDIDS.items():
        if edid == v[1]:
            return k
    raise RuntimeError(f"Couldn't find name associated with EDID: {edid!r}")

@functools.lru_cache(maxsize=8)
def edid_from_name(name: str) -> EDID:
    try:
        return EDIDS[name][1]
    except KeyError:
        raise RuntimeError(f"Couldn't find EDID associated with name {name!r}")

@functools.lru_cache(maxsize=8)
def refresh_rate_from_edid(edid: EDID | str) -> RefreshRate:
    for k, v in EDIDS.items():
        if edid == v[1]:
            return RefreshRate(v[0])
    raise RuntimeError("Couldn't find refresh rate...")

@dataclass
class Monitor:

    name: str
    identifier: str
    dp: str
    connector_type: str
    connector_num: str
    refresh_rate: RefreshRate
    resx: str
    resy: str

    @classmethod
    def from_xrandr_tup(cls, xrandr_tup: tuple[str, str, EDID, str, str, str, str]) -> Monitor:
        dp_, id_, edid_, conn_type, conn_num, resx_, resy_ = xrandr_tup
        edid_ = edid_.replace('\n', '').replace('\t', '').replace(' ', '')
        return Monitor(
            name=name_from_edid(edid=edid_),
            identifier=id_,
            dp=dp_,
            connector_type=conn_type,
            connector_num=conn_num,
            refresh_rate=refresh_rate_from_edid(edid=edid_),
            resx=resx_,
            resy=resy_,
        )

    @property
    def edid(self, ) -> EDID:
        return edid_from_name(name=self.name)

    def model(self) -> str:
        return self.name.split(' ')[-1]

    def __post_init__(self):
        self.refresh_rate = RefreshRate(f"{float(self.refresh_rate):.2f}")


def monitors_from_xrandr_output() -> list[Monitor]:

    xrandr_raw: str = subprocess.check_output(["xrandr", "--verbose"]).decode('utf8')

    monitor_rex = REX_XRANDR.findall(xrandr_raw)

    monitors = []
    for monitor_tup in monitor_rex:
        monitors.append(Monitor.from_xrandr_tup(monitor_tup))
    return monitors

class MonitorConfigTup(NamedTuple):
    config: MonitorLocations
    primary_mon: str

def parse_args(args) -> MonitorConfigTup:
    primary: str = 'MID'
    try:
        config_keyw = args[0]

        if len(args) >= 2:
            primary = args[1]

    except IndexError:
        config_keyw = DEFAULT_CONFIG
        warnings.warn(f"No monitor config provided; defaulting to {config_keyw!r}", UserWarning)

    try:
        config = CONFIGS[config_keyw]
    except KeyError as err:
        raise KeyError(f"Unknown config keyword/ID: {config_keyw!r}.") from err

    return MonitorConfigTup(config, primary)

def get_xrandr_cmd(
        mon_locs: MonitorLocations,
        primary: str,
        monitors: list[Monitor],
        print_cmd: bool = True,
) -> list[str]:

    cmd: list[str] = ['xrandr']

    mon_l, mon_m, mon_r, mon_t = [None] * 4
    if mon_locs[1][0]:
        mon_l = [mon for mon in monitors if mon.name == PHYSICAL_LOCS['LEFT']]
        mon_l = mon_l[0] if mon_l else None

    if mon_locs[1][1]:
        mon_m = [mon for mon in monitors if mon.name == PHYSICAL_LOCS['MID']]
        mon_m = mon_m[0] if mon_m else None

    if mon_locs[1][2]:
        mon_r = [mon for mon in monitors if mon.name == PHYSICAL_LOCS['RIGHT']]
        mon_r = mon_r[0] if mon_r else None

    if mon_locs[0][1]:
        mon_t = [mon for mon in monitors if mon.name == PHYSICAL_LOCS['TOP']]
        mon_t = mon_t[0] if mon_t else None

    mon_t_offset_y = 0
    if mon_m and mon_t:
        mon_t_offset_y = -int(mon_t.resy)
    elif any((mon_l, mon_r)):
        mon_t_offset_y = max(
            -int(mon_l.resy) if mon_l is not None else 0,
            -int(mon_r.resy) if mon_r is not None else 0
        )

    cur_x_offset = 0
    mon_t_offset_x = 0

    for iimon, (mon, loc) in enumerate(zip([mon_l, mon_m, mon_r], ['LEFT', 'MID', 'RIGHT'])):
        if (mon_m and iimon == 1) or not mon_m:
            if not mon_m:
                if mon_l:
                    mon_t_offset_x = int(mon_l.resx)
                elif mon_r:
                    mon_t_offset_x = int(mon_r.resx)
                else:
                    mon_t_offset_x = 0
            else:
                mon_t_offset_x = int(cur_x_offset)

        if mon is None:
            continue

        # Offset to make screens at sides align with bottom of neighboring heigher screen.
        flanking_mon_offset_y = 0
        if mon_m:
            if loc == 'LEFT' and mon_l:
                flanking_mon_offset_y += round(int(mon_m.resy) - int(mon_l.resy))
            elif loc == 'RIGHT' and mon_r:
                flanking_mon_offset_y += round(int(mon_m.resy) - int(mon_r.resy))


        cur_moncmd = [
            "--output", mon.dp,
            "--mode", f"{mon.resx}x{mon.resy}",
            "--pos", f"{cur_x_offset}x{flanking_mon_offset_y}",
            "--rotate", "normal",
            "-r", f"{mon.refresh_rate!r}",
        ]
        if loc == primary:
            cur_moncmd.append("--primary")
        cur_x_offset += int(mon.resx)
        cmd.extend(cur_moncmd)

    if mon_t:
        if mon_m:
            # Align top monitor to middle of top screen.
            mon_t_offset_x += round((int(mon_m.resx) - int(mon_t.resx)) / 2)

        cur_moncmd = [
            "--output", mon_t.dp,
            "--mode", f"{mon_t.resx}x{mon_t.resy}",
            "--pos", f"{mon_t_offset_x}x{mon_t_offset_y}",
            "--rotate", "normal",
            "-r", f"{mon_t.refresh_rate!r}",
        ]
        cmd.extend(cur_moncmd)

    for loc1, mon in zip(('LEFT', 'MID', 'RIGHT', 'TOP'), (mon_l, mon_m, mon_r, mon_t)):
        if mon is None:
            try:
                new_mon = [mon for mon in monitors if mon.name == PHYSICAL_LOCS[loc1]][0]
            except IndexError:
                print(f"Monitor at position {loc1} is probably unplugged. Skipping its command.")
                continue
            cmd.extend(["--output", new_mon.dp, "--off"])


    if print_cmd:
        print(escape(f"Command: \n\t {' '.join(cmd)}"))
    if len(cmd) <= 1:
        raise RuntimeError("No xrandr args could be generated somehow...")

    return cmd


def main(args: list[str]):
    mon_locs_, primary_ = parse_args(args=args)

    monitors_: list[Monitor] = monitors_from_xrandr_output()

    try:
        print(f"Changing monitor setup to {sys.argv[1]!r}")
    except IndexError:
        print(f"Changing monitor setup to default: {DEFAULT_CONFIG!r}")

    cmdlist = get_xrandr_cmd(mon_locs=mon_locs_, primary=primary_, monitors=monitors_)
    cmd = ' '.join(cmdlist)
    os.system(cmd)

def run_tests(wait_between_sec: int = 8):
    start_time = time.time()

    #PRIMARIES = ('LEFT', 'MID', 'RIGHT', 'TOP')
    PRIMARIES = ('MID',)



    monitors_: list[Monitor] = monitors_from_xrandr_output()
    print(f"Testing {len(CONFIGS)} configurations with {len(PRIMARIES)} primaries...")
    print(f"Total no. configurations: {len(CONFIGS)*len(PRIMARIES)}")
    print(f"Tests will take approx. {len(CONFIGS) * len(PRIMARIES) * wait_between_sec} seconds...")
    time.sleep(4)
    print("waiting 4 seconds before starting test...")
    for iip, primary_ in enumerate(PRIMARIES):
        for iia, arg in enumerate(CONFIGS):
            if arg != 'off':
                args = [arg, primary_]
                mon_locs_, primary_ = parse_args(args)
                print(f"Testing config {arg!r} with primary {primary_!r}...")
            else:
                args = [arg]
                mon_locs_, primary_ = parse_args(args)
                print(f"Testing config {arg!r} with NO primary...")
            try:
                print(f"Changing monitor setup to {arg!r}")
            except IndexError:
                print(f"Changing monitor setup to default: {DEFAULT_CONFIG!r}")

            cmdlist = get_xrandr_cmd(mon_locs=mon_locs_, primary=primary_, monitors=monitors_)
            cmd = ' '.join(cmdlist)
            os.system(cmd)

            if (iip+1)*(iia+1) != len(PRIMARIES) * len(CONFIGS):
                print(f"Waiting {wait_between_sec} seconds...")
                time.sleep(wait_between_sec)
    print("Tests completed! Returning to state with all monitors enabled!")
    args = ["all", "MID"]
    mon_locs_, primary_ = parse_args(args)
    cmdlist = get_xrandr_cmd(mon_locs=mon_locs_, primary=primary_, monitors=monitors_)
    cmd = ' '.join(cmdlist)
    os.system(cmd)

    print(f"All tests including reset to 'all' state took {time.time()-start_time} seconds.")

if __name__ == '__main__':
    DEFAULT_TEST_WAIT: int = 15
    try:
        if sys.argv[1].lower() in ('help', '--help', '-help'):
            print("Tool for changing monitor configurations for 4 specific monitors in a .:. configuration.\n\n")
            print("monitor [CONFIG_KEY] [PRIMARY_MONITOR]\n\n")
            print("Config key values:\n")
            print(f"\t {', '.join((f'{x!r}' for x in CONFIGS))}")
            print("\nPRIMARY_MONITOR values:\n ")
            print(f"\t {', '.join((f'{x!r}' for x in ('LEFT', 'MID', 'RIGHT', 'TOP')))}\n")

            sys.exit(0)
        elif sys.argv[1] == 'test':
            if len(sys.argv) > 2:
                wait_between_sec_ = int(sys.argv[2])
            else:
                wait_between_sec_ = DEFAULT_TEST_WAIT
            run_tests(wait_between_sec=wait_between_sec_)
        else:
            main(args=sys.argv[1:])
    except KeyboardInterrupt:
        # Reset monitors to mid
        main(args=['mid', 'MID'])
    except Exception as err:
        print(f"ERROR OCCURED; attempting restore to 'mid'... (error: {str(err)!r})")
        try:
            main(args=['mid', 'MID'])
        except Exception as err:
            print(f"ERROR occured trying to restore; exiting...")
            raise Exception(f"ERROR occured trying to restore; exiting...") from err