Update README.md
[rainbowstream.git] / rainbowstream / rainbow.py
CommitLineData
91476ec3
O
1"""
2Colorful user's timeline stream
3"""
4
5from __future__ import print_function
6
7import os, os.path, argparse, random
8
9from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
10from twitter.oauth import OAuth, read_token_file
11from twitter.util import printNicely
12from dateutil import parser
13from pyfiglet import figlet_format
14
15from colors import *
16from config import *
17
18def 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
30def 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 userlen = len(name + screen_name) + 3
47 clocklen = len(time) + 2
48
49 # Draw rainbow
50 terminalrows, terminalcolumns = os.popen('stty size', 'r').read().split()
51 line1 = u"{u:>{uw}} {t:>{tw}}".format(
52 u = user,
53 uw = len(user) + 2,
54 t = clock,
55 tw = len(clock) + int(terminalcolumns) - userlen - 4 - clocklen
56 )
57
58 line2 = ' ' + tweet
59 line3 = '\n'
60
61 return line1, line2, line3
62
63
64def parse_arguments():
65 """
66 Parse the arguments
67 """
68
69 parser = argparse.ArgumentParser(description=__doc__ or "")
70
71 parser.add_argument('-to', '--timeout', help='Timeout for the stream (seconds).')
72 parser.add_argument('-ht', '--heartbeat-timeout', help='Set heartbeat timeout.', default=90)
73 parser.add_argument('-nb', '--no-block', action='store_true', help='Set stream to non-blocking.')
74 parser.add_argument('-tt', '--track-keywords', help='Search the stream for specific text.')
75 return parser.parse_args()
76
77
78def main():
79 args = parse_arguments()
80
81 # The Logo
82 asciiart()
83
84 # When using twitter stream you must authorize.
85 oauth_filename = os.environ.get('HOME', os.environ.get('USERPROFILE', '')) + os.sep + '.twitter_oauth'
86 oauth_token, oauth_token_secret = read_token_file(oauth_filename)
87 auth = OAuth(oauth_token, oauth_token_secret, CONSUMER_KEY, CONSUMER_SECRET)
88
89 # These arguments are optional:
90 stream_args = dict(
91 timeout=args.timeout,
92 block=not args.no_block,
93 heartbeat_timeout=args.heartbeat_timeout)
94
95 # Track keyword
96 query_args = dict()
97 if args.track_keywords:
98 query_args['track'] = args.track_keywords
99
100 # Get stream
101 stream = TwitterStream(auth=auth, domain='userstream.twitter.com', **stream_args)
102 tweet_iter = stream.user(**query_args)
103
104 # Iterate over the sample stream.
105 for tweet in tweet_iter:
106 if tweet is None:
107 printNicely("-- None --")
108 elif tweet is Timeout:
109 printNicely("-- Timeout --")
110 elif tweet is HeartbeatTimeout:
111 printNicely("-- Heartbeat Timeout --")
112 elif tweet is Hangup:
113 printNicely("-- Hangup --")
114 elif tweet.get('text'):
115 line1, line2, line3 = draw(t = tweet)
116 printNicely(line1)
117 printNicely(line2)
118 printNicely(line3)
119
120if __name__ == '__main__':
121 main()