1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254

#!/usr/bin/env python3
"""
重置当前工作区所有列的宽度为等比模式，修复鼠标拖拽导致的列宽溢出和左侧间隙。

=== 问题背景 ===

Niri 中鼠标右键拖拽窗口边缘调整列宽时，会把列宽设为固定像素值（SetFixed）。
之后无论新增窗口、切换显示器，该列都死守这个像素值不变。
当两列固定宽度之和超出屏幕 → 溢出，小于屏幕 → 留空。
自动贴砖脚本（niri_tile_to_n.py）的 consume/expel 操作也可能在左侧产生空列。

=== 解决原理 ===

通过 niri IPC 逐列聚焦，用 SetProportion 把固定像素宽度恢复成等比分配。
之后 Niri 会动态分配空间，不再溢出或留空。

=== 关于视口滚动的处理 ===

因脚本需要逐列 FocusWindow 来 SetColumnWidth，会导致 Niri 的滚动视口
跟随焦点偏移，最后停留在最后一列上。脚本末尾分三步修复：
  1. FocusColumn(0)    → 把视口拉到第一列（焦点暂移）
  2. CenterVisibleColumns → 居中所有可见列（不改变焦点）
  3. FocusWindow(orig)  → 恢复原始焦点（窗口已在视口内，不触发额外滚动）

=== 用法 ===

  niri_fix_columns.py                      # 所有列平分宽度
  niri_fix_columns.py 2/3_1/3              # 奇数列 2/3，偶数列 1/3
  niri_fix_columns.py --pattern 1/2        # 所有列设成 50%
  niri_fix_columns.py 1/3_1/3_1/3          # 三列各 1/3

快捷键绑定（加到 niri config.kdl 的 binds {} 里）：
  Mod+Shift+R { spawn-sh \"$REPO_DIR/script/niri_fix_columns.py 2/3_1/3\"; }
"""

import json
import os
import socket
import sys
import subprocess
from collections import defaultdict


def notify(msg):
    """通过桌面通知向用户反馈结果，桌面环境不可用时 fallback 到 stderr"""
    try:
        subprocess.run(["notify-send", "-a", "niri-fix-columns", msg], timeout=3)
    except Exception:
        print(msg, file=sys.stderr)


def get_socket_path():
    """获取 niri IPC socket 路径，从环境变量 NIRI_SOCKET 读取"""
    path = os.environ.get("NIRI_SOCKET")
    if not path:
        notify("❌ 找不到 NIRI_SOCKET，确定在 niri 里面运行的？")
        sys.exit(1)
    return path


def send_and_recv(sock, msg):
    """
    发送一条 JSON 消息到 niri IPC，读取一行 JSON 响应。

    niri IPC 是 Unix Socket 上的换行分隔 JSON 协议（NDJSON），
    每条消息以换行符结尾，响应也是单行 JSON。
    """
    sock.sendall((json.dumps(msg) + "\n").encode())
    buf = b""
    while True:
        chunk = sock.recv(4096)
        if not chunk:
            break
        buf += chunk
        if b"\n" in buf:
            break
    line = buf.split(b"\n", 1)[0]
    return json.loads(line.decode())


def request(sock, req_name):
    """
    发送简单请求（如 "Windows" / "Workspaces"），返回 Ok 包裹的数据内容。

    niri 请求的响应格式为 {"Ok": {"<请求名>": <数据>}} 或 {"Err": "..."}。
    例如：
      请求 "Windows"     → {"Ok": {"Windows": [...]}}
      请求 "Workspaces"  → {"Ok": {"Workspaces": [...]}}
      请求 "FocusedWindow" → {"Ok": {"FocusedWindow": {...}}}
    """
    resp = send_and_recv(sock, req_name)
    if "Err" in resp:
        notify(f"❌ niri 返回错误：{resp['Err']}")
        sys.exit(1)
    return resp["Ok"][req_name]


def do_action(sock, action_name, **kwargs):
    """
    执行一个 niri action。

    构建的 IPC 消息格式：
      {"Action": {"<动作名>": {<参数>}}}

    例如 FocusWindow(id=123) 发送：
      {"Action": {"FocusWindow": {"id": 123}}}

    例如 SetColumnWidth(change={"SetProportion": 50.0}) 发送：
      {"Action": {"SetColumnWidth": {"change": {"SetProportion": 50.0}}}}
    """
    msg = {"Action": {action_name: kwargs}}
    send_and_recv(sock, msg)


def parse_proportion_pattern(pattern):
    """
    解析比例模式字符串，返回每列比例值列表（值域 0~1）。

    支持的模式：
      "even" / "1/2"          → [50%, 50%]
      "2/3_1/3"              → [66.7%, 33.3%]
      "1/3_2/3"              → [33.3%, 66.7%]
      "1/3_1/3_1/3"          → [33.3%, 33.3%, 33.3%]
      "60_40"                → [60%, 40%]（自定义百分比）
    多于预设列数时，比例循环使用（奇偶列模式会自动延续）。
    """
    known = {
        "even": [0.5, 0.5],
        "1/2": [0.5, 0.5],
        "2/3_1/3": [2.0 / 3, 1.0 / 3],
        "1/3_2/3": [1.0 / 3, 2.0 / 3],
        "1/3_1/3_1/3": [1.0 / 3, 1.0 / 3, 1.0 / 3],
        "1/4_3/4": [0.25, 0.75],
        "3/4_1/4": [0.75, 0.25],
    }
    if pattern in known:
        return known[pattern]

    # 尝试解析自定义格式
    try:
        parts = pattern.split("_")
        values = []
        for p in parts:
            if "/" in p:
                a, b = p.split("/")
                values.append(float(a) / float(b))
            else:
                values.append(float(p) / 100.0)
        if values:
            return values
    except Exception:
        pass

    notify(f"❌ 无法解析模式 '{pattern}'")
    sys.exit(1)


def main():
    # ======================== 参数解析 ========================
    pattern = "even"
    if len(sys.argv) > 1:
        if sys.argv[1] == "--pattern" and len(sys.argv) > 2:
            pattern = sys.argv[2]
        else:
            pattern = sys.argv[1]

    proportions = parse_proportion_pattern(pattern)

    # ======================== 连接 niri IPC ========================
    skt_path = get_socket_path()
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    sock.connect(skt_path)

    # ======================== 找出当前工作区 ========================
    workspaces = request(sock, "Workspaces")
    focused_wspace_id = None
    for ws in workspaces:
        if ws.get("is_focused"):
            focused_wspace_id = ws["id"]
            break
    if focused_wspace_id is None:
        notify("❌ 找不到当前聚焦的工作区")
        sys.exit(1)

    # ======================== 按列分组当前工作区的窗口 ========================
    # 每个窗口的 layout.pos_in_scrolling_layout = [列索引, 行索引]
    # 例如 [0, 0] = 第一列第一行, [1, 0] = 第二列第一行
    all_windows = request(sock, "Windows")
    columns = defaultdict(list)
    for win in all_windows:
        if win.get("workspace_id") != focused_wspace_id:
            continue
        if win.get("is_floating"):
            continue
        layout = win.get("layout", {})
        pos = layout.get("pos_in_scrolling_layout")
        if pos is None:
            continue
        col_idx = pos[0]
        columns[col_idx].append(win)

    if not columns:
        notify("ℹ️ 当前工作区没有平铺窗口")
        sys.exit(0)

    sorted_cols = sorted(columns.keys())

    # ======================== 记住原始焦点 ========================
    # 脚本最后会恢复焦点，不改变用户的操作位置
    focused = request(sock, "FocusedWindow")
    orig_focus_id = focused.get("id")

    # ======================== 逐列设置等比宽度 ========================
    # SetColumnWidth 只对当前聚焦的列生效，所以需要逐列 FocusWindow
    for i, col_idx in enumerate(sorted_cols):
        first_win = columns[col_idx][0]
        target_id = first_win["id"]

        # 聚焦该列的第一个窗口
        do_action(sock, "FocusWindow", id=target_id)

        # 取该列的目标比例，循环使用
        # （例如 pattern=2/3_1/3 时：列0=66.7%, 列1=33.3%, 列2=66.7%, ...）
        pct = proportions[i % len(proportions)] * 100

        # 设回等比宽度（关键！把固定像素值恢复成动态比例）
        do_action(sock, "SetColumnWidth", change={"SetProportion": pct})

    # ======================== 恢复焦点，同时保持视口正确 ========================
    # 逐列 FocusWindow 会导致 niri 的滚动视口跟随焦点偏移。
    # 如果最后聚焦的是第二列，视口就卡在第二列，第一列滚到屏幕左边外面去了。
    #
    # 修复策略（三步，无视觉闪烁）：
    #   1. FocusColumn(0)     → 把视口拉到第一列，确保所有列在屏幕内
    #   2. CenterVisibleColumns → 居中可见列（不改变焦点），消除偏移
    #   3. FocusWindow(orig)  → 恢复原始焦点，此时窗口已在可见区域，不触发滚动

    do_action(sock, "FocusColumn", index=0)           # 1. 视口拉到列0（焦点暂移）
    do_action(sock, "CenterVisibleColumns")            # 2. 居中，列填满屏幕
    if orig_focus_id is not None:
        do_action(sock, "FocusWindow", id=orig_focus_id)   # 3. 恢复焦点（不滚动）

    sock.close()

    # ======================== 桌面通知反馈 ========================
    parts = []
    for i in range(len(sorted_cols)):
        p = proportions[i % len(proportions)]
        parts.append(f"列{i+1} {p*100:.0f}%")
    notify(f"✓ 已重置 {len(sorted_cols)} 列")


if __name__ == "__main__":
    main()