126 lines
4.2 KiB
Python
126 lines
4.2 KiB
Python
"""Tree-walking helpers for thread spine and branch construction."""
|
|
from __future__ import annotations
|
|
|
|
from .models import ThreadData, PostData
|
|
|
|
|
|
def by_id(thread: ThreadData) -> dict[str, PostData]:
|
|
return {p.post_id: p for p in thread.posts}
|
|
|
|
|
|
def children_map(thread: ThreadData) -> dict[str, list[str]]:
|
|
ids = by_id(thread)
|
|
kids: dict[str, list[str]] = {p.post_id: [] for p in thread.posts}
|
|
for p in thread.posts:
|
|
if p.reply_to_id and p.reply_to_id in ids:
|
|
kids[p.reply_to_id].append(p.post_id)
|
|
for parent, lst in kids.items():
|
|
lst.sort(key=lambda i: (ids[i].timestamp, ids[i].post_id))
|
|
return kids
|
|
|
|
|
|
def spine_from_tip(thread: ThreadData, tip_id: str) -> list[str]:
|
|
"""Resolve the OP chain from `tip_id`.
|
|
|
|
If `tip_id` is the OP (its `reply_to_id` is `None`), walk DOWN
|
|
through same-handle children only. Stop when the OP has no
|
|
self-reply. Do not promote a foreign-handle child onto the
|
|
spine; those posts are branches. This is the `--id <OP> --tip`
|
|
case. The back-walk would otherwise return a single-element
|
|
chain and drop every self-reply.
|
|
|
|
If `tip_id` has a parent, walk BACK through `reply_to_id` to
|
|
the OP and return the chain in chronological order (OP at
|
|
index 0, tip at last). The walk crosses handle boundaries.
|
|
"""
|
|
ids = by_id(thread)
|
|
if tip_id not in ids:
|
|
raise ValueError(f"tip {tip_id} not in dump")
|
|
if ids[tip_id].reply_to_id is None:
|
|
# Tip is the OP. Walk down its same-handle chain only.
|
|
kids = children_map(thread)
|
|
handle = ids[tip_id].handle
|
|
out = [tip_id]
|
|
cur = tip_id
|
|
while True:
|
|
children = kids.get(cur, [])
|
|
same = [c for c in children if ids[c].handle == handle]
|
|
if not same:
|
|
break
|
|
cur = same[0]
|
|
out.append(cur)
|
|
return out
|
|
chain = [tip_id]
|
|
cur = tip_id
|
|
while True:
|
|
parent_id = ids[cur].reply_to_id
|
|
if not parent_id or parent_id not in ids:
|
|
break
|
|
chain.append(parent_id)
|
|
cur = parent_id
|
|
chain.reverse()
|
|
return chain
|
|
|
|
|
|
def spine_ids(thread: ThreadData) -> list[str]:
|
|
"""Walk the OP's own chain down through `children_map`.
|
|
|
|
At each level, prefer a same-handle child (to continue the OP's
|
|
own chain) over any child. Fall back to the earliest child when
|
|
no same-handle child exists.
|
|
|
|
The OP is the post with `reply_to_id == None`. The returned
|
|
list is the OP's chain: the OP at index 0, then same-author
|
|
replies in chronological order. Cross-author replies to spine
|
|
posts are branches (rendered as separate files by the emit).
|
|
|
|
For a thread where the OP has no self-replies but cross-author
|
|
replies (e.g., a post that attracted only outside replies), the
|
|
walker picks the earliest child to start a chain — that branch
|
|
becomes the only spine entry after the OP, and any other
|
|
children of the OP become branches.
|
|
"""
|
|
ids = by_id(thread)
|
|
kids = children_map(thread)
|
|
ops = [p.post_id for p in thread.posts if p.reply_to_id is None]
|
|
if ops:
|
|
start = ops[0]
|
|
else:
|
|
start = thread.posts[0].post_id
|
|
handle = ids[start].handle
|
|
out = [start]
|
|
cur = start
|
|
while True:
|
|
children = kids.get(cur, [])
|
|
if not children:
|
|
break
|
|
same = [c for c in children if ids[c].handle == handle]
|
|
cur = same[0] if same else children[0]
|
|
out.append(cur)
|
|
return out
|
|
|
|
|
|
def branch_roots(thread: ThreadData, spine: list[str]) -> list[str]:
|
|
spine_set = set(spine)
|
|
ids = by_id(thread)
|
|
roots: list[str] = []
|
|
for p in thread.posts:
|
|
if p.post_id in spine_set:
|
|
continue
|
|
parent_on_spine = p.reply_to_id in spine_set
|
|
parent_missing = p.reply_to_id is not None and p.reply_to_id not in ids
|
|
if parent_on_spine or parent_missing:
|
|
roots.append(p.post_id)
|
|
roots.sort(key=lambda i: (ids[i].timestamp, ids[i].post_id))
|
|
return roots
|
|
|
|
|
|
def descendants(start_id: str, kids: dict[str, list[str]]) -> list[str]:
|
|
out: list[str] = []
|
|
stack = list(reversed(kids.get(start_id, [])))
|
|
while stack:
|
|
n = stack.pop()
|
|
out.append(n)
|
|
stack.extend(reversed(kids.get(n, [])))
|
|
return out
|