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()
|