Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import argparse
- import psycopg2
- import psycopg2.extras
- try:
- import igraph
- except ImportError:
- print("""
- Please install igraph library
- pip install python-igraph cairocffi
- """)
- raise
- def get_data(db_name, user, password, schema, table):
- conn = psycopg2.connect(f"dbname={db_name} user={user} password={password}")
- # create a cursor
- cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
- cur.execute(f'SELECT * FROM {schema}.{table}')
- data = cur.fetchall()
- cur.close()
- conn.close()
- return data
- def create_graph(data):
- g = igraph.Graph(directed=True)
- # Add vertices
- for db_row in data:
- label = f"""{db_row['term_id']}
- {db_row['term_name'][:24]}
- LVL:{db_row['level']} - SO:{db_row['sort_order']}
- L:{db_row['lft']} - R:{db_row['rght']}"""
- g.add_vertex(name=db_row["term_id"], label=label, **db_row)
- # map id to vertex
- vertices = {v["term_id"]: v for v in g.vs}
- # Add edges
- for term_id, v in vertices.items():
- parent_id = v["parent_id"]
- if parent_id is not None:
- parent_v = vertices[parent_id]
- kwds = {"color": get_edge_color(parent_v, v)}
- g.add_edge(parent_v, v, **kwds)
- # check vertices
- for v in g.vs:
- color_left_right(g, v)
- return g
- def get_edge_color(parent_v, v):
- if parent_v["level"] != v["level"] - 1:
- return "red"
- return {0: "blue", 1: "brown"}.get(parent_v["level"], "black")
- def color_left_right(graph, source_vertex):
- if "color" in source_vertex.attributes() and source_vertex["color"] is not None:
- return (source_vertex["sort_order"], source_vertex["true_lft"], source_vertex["true_rght"])
- source_vertex["color"] = "pink"
- children_so_left_right = []
- children = graph.es.select(_source_in=[source_vertex.index])
- for child_edge in children:
- child_v = graph.vs[child_edge.target]
- child_data = color_left_right(graph, child_v)
- children_so_left_right.append(child_data)
- so, left, right = source_vertex["sort_order"], source_vertex["lft"], source_vertex["rght"]
- # in case of no child
- if not children_so_left_right:
- if right != left + 1:
- source_vertex["color"] = "yellow"
- source_vertex["true_lft"] = left
- source_vertex["true_rght"] = right
- return (so, left, right)
- # check sort_order matches left and right
- children_so_left_right.sort()
- prev = children_so_left_right[0]
- for child in children_so_left_right[1:]:
- _, _, p_rght = prev
- _, c_lft, _ = child
- if p_rght + 1 != c_lft:
- source_vertex["color"] = "blue"
- prev = child
- source_vertex["true_lft"] = min(c[1] for c in children_so_left_right) - 1
- source_vertex["true_rght"] = max(c[2] for c in children_so_left_right) + 1
- if source_vertex["true_lft"] != source_vertex["lft"] or \
- source_vertex["true_rght"] != source_vertex["rght"]:
- source_vertex["color"] = "red"
- return so, source_vertex["true_lft"], source_vertex["true_rght"]
- def plot_graph(g, layout="kk"):
- visual_style = {
- "layout": g.layout(layout),
- "edge_width": 2,
- "bbox": (2400, 2400),
- "vertex_size": 100,
- "margin": 160,
- }
- igraph.plot(g, **visual_style)
- help_msg = """
- Plot MPTT table.
- # pip install python-igraph cairocffi
- Legend:
- Vertices:
- pink - OK
- blue - sort order does not match mptt order
- red - issue with children mptt left - right index
- yellow - mptt issue of leaf vertex
- Edges:
- Level 0 -> Level 1 - blue
- Level 1 -> Level 2 - brown
- other - black
- """
- if __name__ == "__main__":
- parser = argparse.ArgumentParser(description=help_msg)
- parser.add_argument('--db', '-d', help='Name of database', default='postgres')
- parser.add_argument('--user', '-u', help='Database user', default='pgadmin')
- parser.add_argument('--password', '-p', help='User password', required=True)
- parser.add_argument('--schema', '-s', help='Schema', default='cmc')
- parser.add_argument('--table', '-t', help='Table', default='analytic_term_tree')
- args = parser.parse_args()
- data = get_data(
- db_name=args.db,
- user=args.user,
- password=args.password,
- schema=args.schema,
- table=args.table
- )
- g = create_graph(data)
- plot_graph(g)
Add Comment
Please, Sign In to add comment