87 lines
2.4 KiB
Python
87 lines
2.4 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
import psycopg2
|
||
import os
|
||
import sys
|
||
from dotenv import load_dotenv
|
||
|
||
|
||
load_dotenv("../../.env", override=True)
|
||
DB_CONFIG = {
|
||
"host": os.getenv("ROOT_DB_HOST", "localhost"),
|
||
"port": os.getenv("ROOT_DB_PORT", "5432"),
|
||
"dbname": os.getenv("ROOT_DB_NAME", "mydb"),
|
||
"user": os.getenv("ROOT_DB_USER", "postgres"),
|
||
"password": os.getenv("ROOT_DB_PASSWD", "postgres"),
|
||
}
|
||
|
||
|
||
def fetch_nodes():
|
||
conn = psycopg2.connect(**DB_CONFIG)
|
||
cur = conn.cursor()
|
||
cur.execute("SELECT id, name, parent_id FROM fst")
|
||
rows = cur.fetchall()
|
||
cur.close()
|
||
conn.close()
|
||
return rows
|
||
|
||
|
||
def build_tree(rows):
|
||
nodes = {}
|
||
for id_, name, parent_id in rows:
|
||
nodes[id_] = {"id": id_, "name": name, "parent_id": parent_id, "children": []}
|
||
roots = []
|
||
for n in nodes.values():
|
||
if n["parent_id"] is None or n["parent_id"] not in nodes:
|
||
roots.append(n)
|
||
else:
|
||
nodes[n["parent_id"]]["children"].append(n)
|
||
|
||
def sort_rec(n):
|
||
n["children"].sort(key=lambda x: x["name"])
|
||
for c in n["children"]:
|
||
sort_rec(c)
|
||
|
||
for r in roots:
|
||
sort_rec(r)
|
||
return roots
|
||
|
||
|
||
def print_subtree(node, prefix="", is_last=True, depth=0):
|
||
# depth==0: 根节点；depth>=1: 各层子节点
|
||
if depth == 0:
|
||
# 根直接打印名字
|
||
print(node["name"])
|
||
else:
|
||
# 非根：打印前置 prefix + 连接符 + 名字
|
||
connector = "└── " if is_last else "├── "
|
||
print(prefix + connector + node["name"])
|
||
# 计算传给下一层孩子的前缀
|
||
if depth == 0:
|
||
# 第一层孩子不缩进，prefix 仍然空
|
||
child_prefix = ""
|
||
else:
|
||
# depth>=1：如果我是本层最后一个，用空格填充，否则保留竖线
|
||
child_prefix = prefix + (" " if is_last else "│ ")
|
||
# 递归打印子节点
|
||
cnt = len(node["children"])
|
||
for idx, child in enumerate(node["children"]):
|
||
last = idx == cnt - 1
|
||
print_subtree(child, child_prefix, last, depth + 1)
|
||
|
||
|
||
def main():
|
||
rows = fetch_nodes()
|
||
if not rows:
|
||
print("fst 表中没有数据。", file=sys.stderr)
|
||
sys.exit(1)
|
||
roots = build_tree(rows)
|
||
for idx, r in enumerate(roots):
|
||
last_root = idx == len(roots) - 1
|
||
print_subtree(r, prefix="", is_last=last_root, depth=0)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|