i3status-wrapper: don't assume initial state is right after startup

1 files changed, 30 insertions, 7 deletions

diff --git a/scripts/desktop/i3status-wrapper b/scripts/desktop/i3status-wrapper
index d741c3b5..97c76034 100755
--- a/scripts/desktop/i3status-wrapper
+++ b/scripts/desktop/i3status-wrapper
@@ -2,7 +2,7 @@
 
 # i3status-wrapper -- wrapper for i3status(1), plus other monitoring
 #
-# Copyright (C) 2019, 2021-2022 Sean Whitton
+# Copyright (C) 2019, 2021-2023 Sean Whitton
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
@@ -80,19 +80,30 @@ my $caffeinated_id;
 open my $events, "-|",
   $wmipc, "-t", "subscribe", "-m", '[ "window", "workspace" ]';
 
+# Determine the initial state -- the WM might just have been reloaded.
+
 my $workspaces = @{ decode_json `$wmipc -t get_workspaces` };
-$workspaces > 1 or wsbuttons("no");
+wsbuttons($workspaces > 1 ? "yes" : "no");
+
+my @trees = decode_json `$wmipc -t get_tree`;
+while (@trees) {
+    for ((shift @trees)->{nodes}->@*) {
+	if (grep $_ eq "caffeinated", $_->{marks}->@*) {
+	    register_caffeinated($_);
+	    last;
+	} else {
+	    unshift @trees, $_;
+	}
+    }
+}
 
 while (my $event = eval { decode_json <$events> } // next) {
     if ($event->{change} eq "mark") {
         if (grep $_ eq "caffeinated", $event->{container}{marks}->@*) {
-            $caffeinated_id = $event->{container}{id};
-            $info{caffeinated_name} = $event->{container}{name};
-            kill USR1 => $i3status;
+	    register_caffeinated($event->{container});
         } elsif ($caffeinated_id
             and $caffeinated_id == $event->{container}{id}) {
-            undef $caffeinated_id, undef $info{caffeinated_name};
-            kill USR1 => $i3status;
+            clear_caffeinated();
         }
     } elsif ($event->{change} eq "init") {
 	++$workspaces == 2 and wsbuttons("yes");
@@ -111,3 +122,15 @@ sub wsbuttons {
     return unless $ENV{XDG_CURRENT_DESKTOP} eq "sway";
     wmipc "bar bar-0 workspace_buttons $_[0]";
 }
+
+sub register_caffeinated {
+    $caffeinated_id = $_[0]->{id};
+    $info{caffeinated_name} = $_[0]->{name};
+    kill USR1 => $i3status;
+}
+
+sub clear_caffeinated {
+    undef $caffeinated_id;
+    undef $info{caffeinated_name};
+    kill USR1 => $i3status;
+}