# -*- coding: utf-8 -*- import os, sys, json import networkx as nx import redis ##### # Seen this before - it just writes a dot file # def write_dot_file(graph,fname): try: nx.drawing.write_dot(graph,fname) except ImportError, e: dot = ['"%s" -> "%s"'%(n1, n2) for n1,n2 in graph.edges()] dotenc = [ item.encode('ascii','replace') for item in dot ] f = open(fname,"w") f.write('digraph {\n%s\n}'%(';\n'.join(dotenc),)) f.close() ##### # # Redis access code # ##### # # Take a profile or user id, generate user id keys, save/update the following list # In Twitter the people you are following are also known as "friends" # def redisUpdateFriendList(r=None, prof=None, uid=None, flist=None): assert prof is not None or uid is not None assert r is not None and flist is not None value_type = "friend_list" if( prof ): user_id_key = "user_id$"+str(prof['id'])+"$"+value_type else: user_id_key = "user_id$"+str(uid)+"$"+value_type # create a list comprehension of follower ID strings follow_list = [ str(i) for i in flist ] # add the list of followers via a user ID r.sadd(user_id_key, follow_list) return # # Take a profile or user id, generate user id keys, save/update the follower list # This is the list of people who are followers of this user # def redisUpdateFollowerList(r=None, prof=None, uid=None, flist=None): assert prof is not None or uid is not None assert r is not None and flist is not None value_type = "follower_list" if( prof ): user_id_key = "user_id$"+str(prof['id'])+"$"+value_type else: user_id_key = "user_id$"+str(uid)+"$"+value_type # create a list comprehension of follower ID strings follow_list = [ str(i) for i in flist ] # add the list of followers via a user ID r.sadd(user_id_key, follow_list) return # # Given a profile or a user id get the users being followed from redis server # def redisGetFriendList(r=None, prof=None, uid=None): assert prof is not None or uid is not None assert r is not None value = [] value_type = "friend_list" if( prof ): user_id_key = "user_id$"+str(prof['id'])+"$"+value_type else: user_id_key = "user_id$"+str(uid)+"$"+value_type mset = r.smembers(user_id_key) if( mset ): mlist = [ elt for elt in mset ] slist = mlist[0].replace('[','').replace(']','').replace("'",'').split(',') value = [ int(elt) for elt in slist ] return value # # Given a profile or a user id get the list of followers from redis server # def redisGetFollowerList(r=None, prof=None, uid=None): assert prof is not None or uid is not None assert r is not None value = [] value_type = "follower_list" if( prof ): user_id_key = "user_id$"+str(prof['id'])+"$"+value_type else: user_id_key = "user_id$"+str(uid)+"$"+value_type mset = r.smembers(user_id_key) if( mset ): mlist = [ elt for elt in mset ] slist = mlist[0].replace('[','').replace(']','').replace("'",'').split(',') value = [ int(elt) for elt in slist ] return value # # Take a profile, generate screen name and user id keys, save the profile # def redisUpdateUserProfile(r=None, prof=None): assert r is not None and prof is not None value_type = "info.json" screen_name_key = "screen_name$"+prof['screen_name']+"$"+value_type user_id_key = "user_id$"+str(prof['id'])+"$"+value_type # add or change the current profile - indexed by screen name and user ID r.set(screen_name_key, json.dumps(prof)) r.set(user_id_key, json.dumps(prof)) return # # Given a screen name or a user id get the profile from redis server # def redisGetUserProfile(r=None, user_id=None, screen_name=None): assert screen_name is not None or user_id is not None assert r is not None value = {} value_type = "info.json" if( screen_name ): key = "screen_name$"+screen_name+"$"+value_type else: key = "user_id$"+str(user_id)+"$"+value_type redis_value = r.get(key) if( redis_value ): value = json.loads(redis_value) return value def buildFriendsGraph(user=None,flist=None,g=None,r=None,level=0): assert user is not None and flist is not None assert g is not None and r is not None edges = 0 ## I had to cut the graph down to get it to render in GraphVis f = flist[:50] for uid in f: #for uid in flist: # get our friend profile p = redisGetUserProfile(r=r,user_id=uid) if( p ): # only add an edge if we have this person edges += 1 g.add_edge(user,p['screen_name']) # if we're going more levels out, get the friends of this person if( level > 0 ): pfriends = redisGetFriendList(r=r,uid=uid) # recurse on those friends only if some are listed if( pfriends ): l = level-1 e = buildFriendsGraph(user=p['screen_name'], flist=pfriends, g=g, r=r, level=l) edges = edges + e return edges def parse_params(argv): user_name = "" pc = 1 while( pc < len(argv) ): param = argv[pc] if( param == "-user"): pc += 1 user_name = argv[pc] pc += 1 return {'user_name':user_name } def usage(prog): print "USAGE: %s -user "%(prog) sys.exit(0) def main(argv): if len(argv) < 3: print "ERROR: Must specify a twitter username" usage(sys.argv[0]) params = parse_params(sys.argv) if( not params['user_name'] ): print "ERROR: Must specify a twitter username" usage(sys.argv[0]) db = redis.Redis() # get this user from our redis db - hopefully we have them u_profile = redisGetUserProfile(r=db,screen_name=params['user_name']) if( not u_profile ): print "Error: don't have \"%s\" in the redis DB"%(params['user_name']) sys.exit(0) g = nx.DiGraph() edges = 0 friend_ids = redisGetFriendList(r=db,prof=u_profile) print "len(friend_ids)",len(friend_ids) print "friend_ids",friend_ids if( friend_ids ): edges = buildFriendsGraph(user=u_profile['screen_name'], flist=friend_ids, g=g, r=db, level=1) print "Inserted %d edges in graph, some may have been duplicates"%(edges) write_dot_file(g,"friend.graph.dot") print "Done" if __name__ == '__main__': main(sys.argv)