networkx for python のお試し

networkxを初めて使用しますので、グラフデータ作成やその描画

# -*- coding: utf-8 -*-
import jaconv
import matplotlib.pyplot as plt
import networkx as nx
import sqlite3

#探索するルートノード
root_nodes = ["2C0？？？","2C1Q1？？？CIZW","2C8？？？45JBZL"]

db_path = 'bom.sqlite'
font    = "MS Gothic" # 日本語を表示する為

def main():
    # DB(sqlite)接続
    db_conn = sqlite3.connect( db_path )
    db_conn.row_factory = dict_factory
    db_cur  = db_conn.cursor()
    
    # ルートノード
    org_mims = root_nodes

    for org_mim in org_mims:
        root_node = select_mim_def(db_cur, org_mim)
        if root_node == None:
            print("fail select_mim() %s" % (org_mim))
            continue

        bom = nx.Graph()
        bom.add_node(1,**root_node) # ルートノード追加

        # 以下のように「data=True」とすることで属性も取得
        last_node = list( bom.nodes(data=True) )[-1]
        expand_bom(db_cur,bom, last_node)       # BOM展開
        draw_bom( bom )                         # BOM描画
    db_cur.close()
    db_conn.close()

def draw_bom( bom ):
    pos = nx.spring_layout(bom, k=0.3)
    # ノードと、そのラベルの表示
    nx.draw_networkx_nodes(bom,
                           pos,
                           node_color='#87cefa',
                           node_size=20)

    node_labels = nx.get_node_attributes(bom, 'label')
    nx.draw_networkx_labels(bom, pos, labels=node_labels,
                            font_size  =10,
                            font_family=font)
    # エッジと、そのラベルの表示
    nx.draw_networkx_edges(bom, pos,
                           arrows=True,
                           arrowstyle='->',
                           edge_color='#aaa')
    edge_labels = nx.get_edge_attributes(bom, 'label')
    
    nx.draw_networkx_edge_labels(bom, pos,
                                 edge_labels=edge_labels,
                                 font_family=font)
    plt.axis('off')
    plt.show()
    
def expand_bom(db_cur, bom, parent_node):

    for child_attr in select_children(db_cur,parent_node):
        child_id = list( bom.nodes )[-1] + 1
        child_attr["label"] = child_attr["code"]+"\n"+child_attr["name"]
        child_attr["label"] = child_attr["label"].replace(" ","\n")
        child_attr["label"] = jaconv.z2h(child_attr["label"],
                                         kana=True,
                                         digit=True,
                                         ascii=True )
        # add_node()やadd_edge()する際のIDはunique
        bom.add_node(child_id,**child_attr)
        bom.add_edge(parent_node[0],
                     child_id,
                     label =child_attr["use_amount"],
                     weight=child_attr["use_amount"] )
        child_node = list( bom.nodes(data=True) )[-1]
        expand_bom(db_cur, bom, child_node)

        
def select_children( db_cur, parent_node ):
    ret_datas = []
    ret_datas += select_mim_children(db_cur, parent_node)
    ret_datas += select_im_children(db_cur, parent_node)
    return ret_datas
    
# sql selectした結果をdict型で取得する為
def dict_factory(cursor, row):
   d = {}
   for idx, col in enumerate(cursor.description):
       d[col[0]] = row[idx]
   return d

def select_im_children(db_cur,parent_node):
    
    sql = """
SELECT mps.m_p_child as code,
       im.name_kanji as name,
       mps.use_amount
FROM mps
JOIN im ON (mps.m_p_child=im.part_num)
WHERE m_p_child_flag='P' AND m_p_parent=?
"""
    sql_vals = (parent_node[1]["code"],)
    db_cur.execute( sql, sql_vals )

    ret_datas = []
    for ret_data in db_cur.fetchall():
        ret_data["type"] = "im"
        ret_datas.append( ret_data )
    return ret_datas

def select_mim_children(db_cur,parent_node):
    
    sql = """
SELECT mps.m_p_child  as code,
       mim.name_kanji as name,
       mps.use_amount
FROM mps
JOIN mim ON (mps.m_p_child=mim.menu_code)
WHERE m_p_child_flag='M' AND m_p_parent=?
"""
    sql_vals = (parent_node[1]["code"],)
    db_cur.execute( sql, sql_vals )

    ret_datas = []
    for ret_data in db_cur.fetchall():
        ret_data["type"] = "mim"
        ret_datas.append( ret_data )
    return ret_datas

def select_mim_def(db_cur, menu_code):
    sql = """
SELECT menu_code as code, name_kanji as name
FROM mim
WHERE menu_code=?
"""
    sql_vals = (menu_code,)
    db_cur.execute( sql, sql_vals )
    ret_data = db_cur.fetchone()
    if not ret_data:
        return
        
    ret_data["type"] = "mim"

    ret_data["label"] = ret_data["code"]+" "+ret_data["name"]
    ret_data["label"] = ret_data["label"].replace(" ","\n")
    ret_data["label"] = jaconv.z2h(ret_data["label"],
                                   kana=True,
                                   digit=True,
                                   ascii=True )
    return ret_data

if __name__ == '__main__':
    main()