Polish pass

hyprland/.config/waybar/scripts/openai-rate.py

@@ -5,6 +5,7 @@ import os
 import sys
 import urllib.error
 import urllib.request
+from argparse import ArgumentParser
 from pathlib import Path
 
 
@@ -107,6 +108,16 @@ def reset_text(seconds):
     return f"{round(hours / 24)}d"
 
 
+def usage_class(used, limit_reached=False):
+    if limit_reached or used >= 95:
+        return "critical"
+    if used >= 90:
+        return "warning"
+    if used >= 75:
+        return "regular"
+    return "good"
+
+
 def format_window(name, window):
     if not isinstance(window, dict):
         return None
@@ -121,7 +132,13 @@ def format_window(name, window):
     tooltip = f"{name} {label}: {used}% used"
     if reset:
         tooltip += f", resets in {reset}"
-    return {"text": f"{left}% {label}", "used": used, "tooltip": tooltip}
+    return {
+        "label": label,
+        "left": left,
+        "text": f"{left}% {label}",
+        "used": used,
+        "tooltip": tooltip,
+    }
 
 
 def credit_text(credits):
@@ -158,11 +175,19 @@ def spend_control_text(spend_control):
     return text
 
 
-def main():
-    token, account_id, source = load_auth()
-    data = fetch_usage(token, account_id)
+def selected_window_output(data, source, selector):
     rate_limit = data.get("rate_limit") or {}
+    window = format_window(selector, rate_limit.get(f"{selector}_window"))
+    if not window:
+        waybar("n/a", "warning", f"No {selector} usage window returned\nsource: {source}")
 
+    klass = usage_class(window["used"], rate_limit.get("limit_reached"))
+    tooltip = [window["tooltip"], f"{window['left']}% remaining", f"source: {source}"]
+    waybar(f"{window['left']}%", klass, "\n".join(tooltip))
+
+
+def combined_output(data, source):
+    rate_limit = data.get("rate_limit") or {}
     windows = [
         format_window("primary", rate_limit.get("primary_window")),
         format_window("secondary", rate_limit.get("secondary_window")),
@@ -178,7 +203,7 @@ def main():
         waybar("AI n/a", "warning", f"No displayable usage data\nsource: {source}")
 
     worst_used = max(window["used"] for window in windows) if windows else 0
-    klass = "critical" if rate_limit.get("limit_reached") or worst_used >= 95 else "warning" if worst_used >= 90 else "good"
+    klass = usage_class(worst_used, rate_limit.get("limit_reached"))
 
     tooltip = [window["tooltip"] for window in windows]
     spend = spend_control_text(data.get("spend_control"))
@@ -191,5 +216,27 @@ def main():
     waybar(f"AI {'  '.join(parts)}", klass, "\n".join(tooltip))
 
 
+def parse_args():
+    parser = ArgumentParser(description="Waybar Codex rate-limit widget")
+    parser.add_argument(
+        "--window",
+        choices=("primary", "secondary", "all"),
+        default="all",
+        help="usage window to display",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    token, account_id, source = load_auth()
+    data = fetch_usage(token, account_id)
+
+    if args.window == "all":
+        combined_output(data, source)
+    else:
+        selected_window_output(data, source, args.window)
+
+
 if __name__ == "__main__":
     main()