Simplify parsing and code

--HG--
extra : convert_revision : svn%3A3ed01bd8-26fb-0310-9e4c-ca1a4053419f/networkx/trunk%401860

1 files changed, 80 insertions, 91 deletions

diff --git a/examples/multigraph/chess_masters.py b/examples/multigraph/chess_masters.py
index ea9c8e4a..f6fe289c 100644
--- a/examples/multigraph/chess_masters.py
+++ b/examples/multigraph/chess_masters.py
@@ -22,7 +22,7 @@ The key statement in chess_pgn_graph below is
 where game_info is a dict describing each game.
 
 """
-#    Copyright (C) 2006 by 
+#    Copyright (C) 2006-2010 by 
 #    Aric Hagberg <hagberg@lanl.gov>
 #    Dan Schult <dschult@colgate.edu>
 #    Pieter Swart <swart@lanl.gov>
@@ -33,13 +33,11 @@ import networkx as nx
 
 # tag names specifying what game info should be 
 # stored in the dict on each digraph edge
-game_details=[
-              "Event", 
+game_details=["Event", 
               "Date", 
               "Result", 
               "ECO",
-              "Site"
-             ]
+              "Site"]
 
 def chess_pgn_graph(pgn_file="chess_masters_WCC.pgn.bz2"):
     """Read chess games in pgn format in pgn_file.
@@ -50,40 +48,31 @@ def chess_pgn_graph(pgn_file="chess_masters_WCC.pgn.bz2"):
     Edges contain game data in a dict.
 
     """
-    try:# use networkx.utils._get_fh to uncompress
-        # pgn file if required
-        datafile=nx.utils._get_fh(pgn_file,mode='rb')
-    except IOError:
-        print("Could not read file %s."%(pgn_file))
-        raise
-    G=nx.MultiDiGraph(weighted=False)
-    game_info={}
-    for line in datafile.read().splitlines():
-        line = line.decode()
-        # check for tag pairs
-        if len(line)>0 and line[0]=='[':
-           line=line[1:-1] # remove extra quotes
-           tag = line.split()[0]
-           value=line[len(tag)+2:-1]
-           if tag=='White':
-              white=value.split(',')[0]
-           elif tag=='Black':
-              black=value.split(',')[0]
-           elif tag in game_details:
-              game_info[tag]=value
+    G=nx.MultiDiGraph()
+    game={}
+    datafile=nx.utils._get_fh(pgn_file)
+    lines = (line.decode().rstrip('\r\n') for line in datafile)
+    for line in lines:
+        if line.startswith('['):
+            tag,value=line[1:-1].split(' ',1)
+            game[str(tag)]=value.strip('"')
+        else:
         # empty line after tag set indicates 
         # we finished reading game info
-        elif len(line)==0:
-             if len(game_info)>0: 
-                 G.add_edge(white, black, **game_info)
-                 game_info={}
+            if game:
+                white=game.pop('White')
+                black=game.pop('Black')
+                G.add_edge(white, black, **game)
+                game={}
     return G
 
 
 if __name__ == '__main__':
     import networkx as nx
 
+
     G=chess_pgn_graph()
+
     ngames=G.number_of_edges()
     nplayers=G.number_of_nodes()
 
@@ -106,70 +95,70 @@ if __name__ == '__main__':
 
     for (white,black,game_info) in G.edges(data=True):
         if game_info['ECO']=='B97':
-           print('%s vs %s' % (white, black))
+           print(white,"vs",black)
            for k,v in list(game_info.items()):
-               print("   %s: %s" % (k, v))
+               print("   ",k,": ",v)
            print("\n")
 
+
     try:
         import matplotlib.pyplot as plt
-
-        plt.rcParams['text.usetex'] = False
-        plt.figure(figsize=(8,8))
-        # make new undirected graph H without multi-edges
-        H=nx.Graph(G)
-        # edge width is proportional number of games played
-        edgewidth=[]
-        for (u,v,d) in H.edges(data=True):
-            edgewidth.append(len(G.get_edge_data(u,v)))
-
-        # node size is proportional to number of games won
-        wins=dict.fromkeys(G.nodes(),0.0)
-        for (u,v,d) in G.edges(data=True):
-            r=d['Result'].split('-')
-            if r[0]=='1':
-                wins[u]+=1.0
-            elif r[0]=='1/2':
-                wins[u]+=0.5
-                wins[v]+=0.5
-            else:
-                wins[v]+=1.0
-
-        print(H.edges(data=True))            
-        A=nx.to_numpy_matrix(H)
-        try:
-            pos=nx.graphviz_layout(H)
-        except:
-            pos=nx.spring_layout(H,iterations=20)
-
-        nx.draw_networkx_edges(H,pos,alpha=0.3,width=edgewidth, edge_color='m')
-        nodesize=[wins[v]*50 for v in H]
-        nx.draw_networkx_nodes(H,pos,node_size=nodesize,node_color='w',alpha=0.4)
-        nx.draw_networkx_edges(H,pos,alpha=0.4,node_size=0,width=1,edge_color='k')
-        nx.draw_networkx_labels(H,pos,fontsize=14)
-        font = {'fontname'   : 'Helvetica',
-                'color'      : 'k',
-                'fontweight' : 'bold',
-                'fontsize'   : 14}
-        plt.title("World Chess Championship Games: 1886 - 1985", font)
-
-        # change font and write text (using data coordinates)
-        font = {'fontname'   : 'Helvetica',
-        'color'      : 'r',
-        'fontweight' : 'bold',
-        'fontsize'   : 14}
-
-        plt.text(0.5, 0.97, "edge width = # games played",
-                 horizontalalignment='center',
-                 transform=plt.gca().transAxes)
-        plt.text(0.5, 0.94,  "node size = # games won",
-                 horizontalalignment='center',
-                 transform=plt.gca().transAxes)
-
-        plt.axis('off')
-        plt.savefig("chess_masters.png",dpi=75)
-        print("Wrote chess_masters.png")
-        plt.show() # display
     except ImportError:
-        pass
-
+        import sys
+        print("Matplotlib needed for drawing. Skipping")
+        sys.exit(0)
+
+    # make new undirected graph H without multi-edges
+    H=nx.Graph(G)
+
+    # edge width is proportional number of games played
+    edgewidth=[]
+    for (u,v,d) in H.edges(data=True):
+        edgewidth.append(len(G.get_edge_data(u,v)))
+
+    # node size is proportional to number of games won
+    wins=dict.fromkeys(G.nodes(),0.0)
+    for (u,v,d) in G.edges(data=True):
+        r=d['Result'].split('-')
+        if r[0]=='1':
+            wins[u]+=1.0
+        elif r[0]=='1/2':
+            wins[u]+=0.5
+            wins[v]+=0.5
+        else:
+            wins[v]+=1.0
+    try:
+        pos=nx.graphviz_layout(H)
+    except:
+        pos=nx.spring_layout(H,iterations=20)
+
+    plt.rcParams['text.usetex'] = False
+    plt.figure(figsize=(8,8))
+    nx.draw_networkx_edges(H,pos,alpha=0.3,width=edgewidth, edge_color='m')
+    nodesize=[wins[v]*50 for v in H]
+    nx.draw_networkx_nodes(H,pos,node_size=nodesize,node_color='w',alpha=0.4)
+    nx.draw_networkx_edges(H,pos,alpha=0.4,node_size=0,width=1,edge_color='k')
+    nx.draw_networkx_labels(H,pos,fontsize=14)
+    font = {'fontname'   : 'Helvetica',
+            'color'      : 'k',
+            'fontweight' : 'bold',
+            'fontsize'   : 14}
+    plt.title("World Chess Championship Games: 1886 - 1985", font)
+
+    # change font and write text (using data coordinates)
+    font = {'fontname'   : 'Helvetica',
+    'color'      : 'r',
+    'fontweight' : 'bold',
+    'fontsize'   : 14}
+
+    plt.text(0.5, 0.97, "edge width = # games played",
+             horizontalalignment='center',
+             transform=plt.gca().transAxes)
+    plt.text(0.5, 0.94,  "node size = # games won",
+             horizontalalignment='center',
+             transform=plt.gca().transAxes)
+
+    plt.axis('off')
+    plt.savefig("chess_masters.png",dpi=75)
+    print("Wrote chess_masters.png")
+    plt.show() # display