Plot MPTT tree

plot_term_tree.py

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)