43fa026999848d1e98c720d5babc2922d6c51644
[rainbowstream.git] / rainbowstream / rainbow.py
1 """
2 Colorful user's timeline stream
3 """
4
5 from __future__ import print_function
6
7 import os, os.path, argparse, random
8
9 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
10 from twitter.oauth import OAuth, read_token_file
11 from twitter.util import printNicely
12 from dateutil import parser
13 from pyfiglet import figlet_format
14
15 from colors import *
16 from config import *
17
18 def asciiart():
19 """
20 Draw the Ascii Art
21 """
22 d = [red, green, yellow, blue, magenta, cyan, white]
23 fi = figlet_format('Rainbow Stream', font='doom')
24 print('\n'.join(
25 [random.choice(d)(i) for i in fi.split('\n')]
26 )
27 )
28
29
30 def draw(t):
31 """
32 Draw the rainbow
33 """
34 # Retrieve tweet
35 text = t['text']
36 screen_name = t['user']['screen_name']
37 name = t['user']['name']
38 created_at = t['created_at']
39 date = parser.parse(created_at)
40 time = date.strftime('%Y/%m/%d %H:%M:%S')
41
42 # Format info
43 user = green(name) + ' ' + yellow('@' + screen_name) + ' '
44 clock = magenta('['+ time + ']')
45 tweet = white(text)
46
47 # Draw rainbow
48 terminalrows, terminalcolumns = os.popen('stty size', 'r').read().split()
49 line1 = u"{u:>{uw}}:".format(
50 u = user,
51 uw = len(user) + 2,
52 )
53 line2 = u"{c:>{cw}}".format(
54 c = clock,
55 cw = len(clock) + 2,
56 )
57 line3 = ' ' + tweet
58 line4 = '\n'
59
60 return line1, line2, line3, line4
61
62
63 def parse_arguments():
64 """
65 Parse the arguments
66 """
67
68 parser = argparse.ArgumentParser(description=__doc__ or "")
69
70 parser.add_argument('-to', '--timeout', help='Timeout for the stream (seconds).')
71 parser.add_argument('-ht', '--heartbeat-timeout', help='Set heartbeat timeout.', default=90)
72 parser.add_argument('-nb', '--no-block', action='store_true', help='Set stream to non-blocking.')
73 parser.add_argument('-tt', '--track-keywords', help='Search the stream for specific text.')
74 return parser.parse_args()
75
76
77 def main():
78 args = parse_arguments()
79
80 # The Logo
81 asciiart()
82
83 # When using twitter stream you must authorize.
84 oauth_filename = os.environ.get('HOME', os.environ.get('USERPROFILE', '')) + os.sep + '.twitter_oauth'
85 oauth_token, oauth_token_secret = read_token_file(oauth_filename)
86 auth = OAuth(oauth_token, oauth_token_secret, CONSUMER_KEY, CONSUMER_SECRET)
87
88 # These arguments are optional:
89 stream_args = dict(
90 timeout=args.timeout,
91 block=not args.no_block,
92 heartbeat_timeout=args.heartbeat_timeout)
93
94 # Track keyword
95 query_args = dict()
96 if args.track_keywords:
97 query_args['track'] = args.track_keywords
98
99 # Get stream
100 stream = TwitterStream(auth=auth, domain='userstream.twitter.com', **stream_args)
101 tweet_iter = stream.user(**query_args)
102
103 # Iterate over the sample stream.
104 for tweet in tweet_iter:
105 if tweet is None:
106 printNicely("-- None --")
107 elif tweet is Timeout:
108 printNicely("-- Timeout --")
109 elif tweet is HeartbeatTimeout:
110 printNicely("-- Heartbeat Timeout --")
111 elif tweet is Hangup:
112 printNicely("-- Hangup --")
113 elif tweet.get('text'):
114 line1, line2, line3, line4 = draw(t = tweet)
115 printNicely(line1)
116 printNicely(line2)
117 printNicely(line3)
118 printNicely(line4)
119
120 if __name__ == '__main__':
121 main()