102 lines
3.2 KiB
Python
102 lines
3.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]:
|
|
"""Walk back from `tip_id` through `reply_to_id` to the OP (the
|
|
post with `reply_to_id == None` or a missing parent). Returns
|
|
all post ids on the linear chain in chronological order: OP at
|
|
index 0, tip at last.
|
|
|
|
The walk crosses handle boundaries. For a chain that alternates
|
|
authors, every post is included. The archive dir for the
|
|
thread is owned by the OP at index 0.
|
|
"""
|
|
ids = by_id(thread)
|
|
if tip_id not in ids:
|
|
raise ValueError(f"tip {tip_id} not in dump")
|
|
chain = [tip_id]
|
|
cur = tip_id
|
|
while True:
|
|
parent = ids[cur].reply_to_id
|
|
if not parent or parent not in ids:
|
|
break
|
|
chain.append(parent)
|
|
cur = parent
|
|
chain.reverse()
|
|
return chain
|
|
|
|
|
|
def spine_ids(thread: ThreadData) -> list[str]:
|
|
"""Walk the linear reply chain from the OP (the post with
|
|
`reply_to_id == None`) down through `children_map`, picking the
|
|
first child at each level (already sorted by `(timestamp,
|
|
post_id)`).
|
|
|
|
The walk crosses handle boundaries. For a chain that alternates
|
|
authors (e.g., rianflo ↔ NOTimothyLottes), every post on the
|
|
chain is included. The archive dir for the thread is owned by
|
|
the OP — the first post in the returned list.
|
|
|
|
If no post has `reply_to_id == None`, the first post in
|
|
`thread.posts` is used as a fallback start.
|
|
"""
|
|
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
|
|
out = [start]
|
|
cur = start
|
|
while True:
|
|
children = kids.get(cur, [])
|
|
if not children:
|
|
break
|
|
cur = 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
|