1 files changed, 49 insertions, 39 deletions

diff --git a/accounts/gkleen@sif/niri/default.nix b/accounts/gkleen@sif/niri/default.nix
index 7254b7ed..3f506d52 100644
--- a/accounts/gkleen@sif/niri/default.nix
+++ b/accounts/gkleen@sif/niri/default.nix
@@ -245,7 +245,7 @@ in {
       Service = {
         Type = "simple";
         Sockets = [ "niri-workspace-history.socket" ];
-        ExecStart = pkgs.writers.writePython3 "niri-workspace-history" {} ''
+        ExecStart = pkgs.writers.writePython3 "niri-workspace-history" { flakeIgnore = ["E501"]; } ''
           import os
           import socket
           import json
@@ -274,7 +274,7 @@ in {
 
               def focus_workspace(output, workspace):
                   with history_lock:
-                      workspace_history[output] = [workspace] + [ws for ws in workspace_history[output] if ws != workspace]  # noqa: E501
+                      workspace_history[output] = [workspace] + [ws for ws in workspace_history[output] if ws != workspace]
                       # print(json.dumps(workspace_history), file=sys.stderr)
 
               sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
@@ -297,14 +297,14 @@ in {
 
           class RequestHandler(StreamRequestHandler):
               def handle(self):
-                  with detaching(TextIOWrapper(self.wfile, encoding='utf-8', write_through=True)) as out:  # noqa: E501
+                  with detaching(TextIOWrapper(self.wfile, encoding='utf-8', write_through=True)) as out:
                       with history_lock:
                           json.dump(workspace_history, out)
 
 
           class Server(ThreadingTCPServer):
               def __init__(self):
-                  ThreadingTCPServer.__init__(self, ("", 8000), RequestHandler, bind_and_activate=False)  # noqa: E501
+                  ThreadingTCPServer.__init__(self, ("", 8000), RequestHandler, bind_and_activate=False)
                   self.socket = socket.fromfd(3, self.address_family, self.socket_type)
 
 
@@ -349,45 +349,55 @@ in {
           outputs = None
           only = {'HDMI-A-1': {'bmr'}, 'eDP-1': {'vid'}}
 
-          with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as cmd_sock, socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
-              cmd_sock.connect(os.environ["NIRI_SOCKET"])
-              with cmd_sock.makefile(mode='w', buffering=1, encoding='utf-8') as fh:
-                  sock.connect(os.environ["NIRI_SOCKET"])
-                  sock.send(b"\"EventStream\"\n")
-                  for line in sock.makefile(buffering=1, encoding='utf-8'):
-                      workspaces = None
-                      if line_json := json.loads(line):
-                          if "WorkspacesChanged" in line_json:
-                              workspaces = line_json["WorkspacesChanged"]["workspaces"]
-
-                      if workspaces is None:
-                          continue
 
-                      old_outputs = outputs
-                      outputs = {ws["output"] for ws in workspaces}
-                      if old_outputs is None:
-                          print("Initial outputs: {}".format(outputs), file=sys.stderr)
-                          continue
+          class Niri(socket.socket):
+              def __init__(self):
+                  super().__init__(socket.AF_UNIX, socket.SOCK_STREAM)
+                  super().connect(os.environ["NIRI_SOCKET"])
+                  self.fh = super().makefile(mode='rw', buffering=1, encoding='utf-8')
+
+              def cmd(self, obj):
+                  print(json.dumps(obj, separators=(',', ':')), flush=True, file=self.fh)
+
+              def event_stream(self):
+                  self.cmd("EventStream")
+                  return self.fh
 
-                      new_outputs = outputs - old_outputs
-                      if not new_outputs:
-                          continue
-                      print("New outputs: {}".format(new_outputs), file=sys.stderr)
 
-                      relevant_workspaces = list(filter(lambda ws: (ws["name"] is not None) or (ws["active_window_id"] is not None), workspaces))
-                      target_output = next(iter(outputs - set(only.keys())))
-                      if not target_output:
+          with Niri() as niri, Niri().event_stream() as niri_stream:
+              for line in niri_stream:
+                  workspaces = None
+                  if line_json := json.loads(line):
+                      if "WorkspacesChanged" in line_json:
+                          workspaces = line_json["WorkspacesChanged"]["workspaces"]
+
+                  if workspaces is None:
+                      continue
+
+                  old_outputs = outputs
+                  outputs = {ws["output"] for ws in workspaces}
+                  if old_outputs is None:
+                      print("Initial outputs: {}".format(outputs), file=sys.stderr)
+                      continue
+
+                  new_outputs = outputs - old_outputs
+                  if not new_outputs:
+                      continue
+                  print("New outputs: {}".format(new_outputs), file=sys.stderr)
+
+                  relevant_workspaces = list(filter(lambda ws: (ws["name"] is not None) or (ws["active_window_id"] is not None), workspaces))
+                  target_output = next(iter(outputs - set(only.keys())))
+                  if not target_output:
+                      continue
+                  for ws in relevant_workspaces:
+                      ws_ident = ws["name"] if ws["name"] is not None else (ws["output"], ws["idx"])
+                      if ws["output"] not in set(only.keys()):
+                          continue
+                      if ws_ident in only[ws["output"]]:
                           continue
-                      for ws in relevant_workspaces:
-                          ws_ident = ws["name"] if ws["name"] is not None else (ws["output"], ws["idx"])
-                          if ws["output"] not in set(only.keys()):
-                              continue
-                          if ws_ident in only[ws["output"]]:
-                              continue
-
-                          print("{} -> {}".format(ws_ident, target_output), file=sys.stderr)
-                          cmd = json.dumps({"Action": {"MoveWorkspaceToMonitor": {"reference": {"Id": ws["id"]}, "output": target_output}}}, separators=(',', ':'))
-                          print(cmd, flush=True, file=fh)
+
+                      print("{} -> {}".format(ws_ident, target_output), file=sys.stderr)
+                      niri.cmd({"Action": {"MoveWorkspaceToMonitor": {"reference": {"Id": ws["id"]}, "output": target_output}}})
         '';
         Restart = "on-failure";
         RestartSec = 10;