9ff0e3971e9f04a9d1ecca0823f35d884fbba681
[rainbowstream.git] / rainbowstream / rainbow.py
1 """
2 Colorful user's timeline stream
3 """
4
5 from __future__ import print_function
6 from multiprocessing import Process
7
8 import os
9 import os.path
10 import argparse
11 import sys
12 import time
13 import signal
14
15 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
16 from twitter.api import *
17 from twitter.oauth import OAuth, read_token_file
18 from twitter.oauth_dance import oauth_dance
19 from twitter.util import printNicely
20 from dateutil import parser
21
22 from .colors import *
23 from .config import *
24
25 g = {}
26
27
28 def draw(t, keyword=None):
29 """
30 Draw the rainbow
31 """
32 # Retrieve tweet
33 text = t['text']
34 screen_name = t['user']['screen_name']
35 name = t['user']['name']
36 created_at = t['created_at']
37 date = parser.parse(created_at)
38 time = date.strftime('%Y/%m/%d %H:%M:%S')
39
40 # Format info
41 user = cycle_color(name) + grey(' ' + '@' + screen_name + ' ')
42 clock = grey('[' + time + ']')
43 tweet = text.split()
44 # Highlight RT
45 tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
46 # Highlight screen_name
47 tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
48 # Highlight link
49 tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
50 # Highlight search keyword
51 if keyword:
52 tweet = map(
53 lambda x: on_yellow(x) if
54 ''.join(c for c in x if c.isalnum()).lower() == keyword.lower()
55 else x,
56 tweet
57 )
58 tweet = ' '.join(tweet)
59
60 # Draw rainbow
61 line1 = u"{u:>{uw}}:".format(
62 u=user,
63 uw=len(user) + 2,
64 )
65 line2 = u"{c:>{cw}}".format(
66 c=clock,
67 cw=len(clock) + 2,
68 )
69 line3 = ' ' + tweet
70 line4 = '\n'
71
72 printNicely(line1)
73 printNicely(line2)
74 printNicely(line3)
75 printNicely(line4)
76
77
78 def parse_arguments():
79 """
80 Parse the arguments
81 """
82 parser = argparse.ArgumentParser(description=__doc__ or "")
83 parser.add_argument(
84 '-to',
85 '--timeout',
86 help='Timeout for the stream (seconds).')
87 parser.add_argument(
88 '-ht',
89 '--heartbeat-timeout',
90 help='Set heartbeat timeout.',
91 default=90)
92 parser.add_argument(
93 '-nb',
94 '--no-block',
95 action='store_true',
96 help='Set stream to non-blocking.')
97 parser.add_argument(
98 '-tt',
99 '--track-keywords',
100 help='Search the stream for specific text.')
101 return parser.parse_args()
102
103
104 def authen():
105 """
106 authenticate with Twitter OAuth
107 """
108 # When using rainbow stream you must authorize.
109 twitter_credential = os.environ.get(
110 'HOME',
111 os.environ.get(
112 'USERPROFILE',
113 '')) + os.sep + '.rainbow_oauth'
114 if not os.path.exists(twitter_credential):
115 oauth_dance("Rainbow Stream",
116 CONSUMER_KEY,
117 CONSUMER_SECRET,
118 twitter_credential)
119 oauth_token, oauth_token_secret = read_token_file(twitter_credential)
120 return OAuth(
121 oauth_token,
122 oauth_token_secret,
123 CONSUMER_KEY,
124 CONSUMER_SECRET)
125
126
127 def get_decorated_name():
128 """
129 Beginning of every line
130 """
131 t = Twitter(auth=authen())
132 name = '@' + t.statuses.user_timeline()[-1]['user']['screen_name']
133 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
134
135
136 def tweet():
137 """
138 Authen and tweet
139 """
140 t = Twitter(auth=authen())
141 t.statuses.update(status=g['stuff'])
142
143
144 def search():
145 """
146 Authen and search
147 """
148 t = Twitter(auth=authen())
149 rel = t.search.tweets(q='#' + g['stuff'])['statuses']
150 h, w = os.popen('stty size', 'r').read().split()
151
152 printNicely(grey('*' * int(w) + '\n'))
153 print('Newest', SEARCH_MAX_RECORD, 'tweet: \n')
154 for i in xrange(5):
155 draw(t=rel[i], keyword=g['stuff'].strip())
156 printNicely(grey('*' * int(w) + '\n'))
157
158
159 def friend():
160 """
161 List of friend (following)
162 """
163 t = Twitter(auth=authen())
164 g['friends'] = t.friends.ids()['ids']
165 for i in g['friends']:
166 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
167 user = cycle_color('@' + screen_name)
168 print(user, end=' ')
169 print('\n')
170
171
172 def follower():
173 """
174 List of follower
175 """
176 t = Twitter(auth=authen())
177 g['followers'] = t.followers.ids()['ids']
178 for i in g['followers']:
179 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
180 user = cycle_color('@' + screen_name)
181 print(user, end=' ')
182 print('\n')
183
184
185 def help():
186 """
187 Print help
188 """
189 usage = '''
190 Hi boss! I'm ready to serve you right now!
191 ----------------------------------------------------
192 "tweet" at the beginning will tweet immediately
193 "s" and follow by any word will search and return 5 newest tweet
194 "fr" will list out your following people
195 "fl" will list out your followers
196 "h" will print this help once again
197 "c" will clear the terminal
198 "q" will exit
199 ----------------------------------------------------
200 Have fun and hang tight!
201 '''
202 printNicely(usage)
203 sys.stdout.write(g['decorated_name'])
204
205
206 def clear():
207 """
208 Exit all
209 """
210 os.system('clear')
211
212
213 def quit():
214 """
215 Exit all
216 """
217 os.kill(g['stream_pid'], signal.SIGKILL)
218 sys.exit()
219
220
221 def process(cmd):
222 """
223 Process switch by start of line
224 """
225 return {
226 'tweet' : tweet,
227 's' : search,
228 'fr' : friend,
229 'fl' : follower,
230 'h' : help,
231 'c' : clear,
232 'q' : quit,
233 }.get(cmd, lambda: sys.stdout.write(g['decorated_name']))
234
235
236 def listen(stdin):
237 """
238 Listen to user's input
239 """
240 for line in iter(stdin.readline, ''):
241 try:
242 cmd = line.split()[0]
243 except:
244 cmd = ''
245 # Save cmd to global variable and call process
246 g['stuff'] = ' '.join(line.split()[1:])
247 process(cmd)()
248 stdin.close()
249
250
251 def stream():
252 """
253 Track the stream
254 """
255 args = parse_arguments()
256
257 # The Logo
258 ascii_art()
259 # These arguments are optional:
260 stream_args = dict(
261 timeout=args.timeout,
262 block=not args.no_block,
263 heartbeat_timeout=args.heartbeat_timeout)
264
265 # Track keyword
266 query_args = dict()
267 if args.track_keywords:
268 query_args['track'] = args.track_keywords
269
270 # Get stream
271 stream = TwitterStream(
272 auth=authen(),
273 domain='userstream.twitter.com',
274 **stream_args)
275 tweet_iter = stream.user(**query_args)
276
277 # Iterate over the sample stream.
278 for tweet in tweet_iter:
279 if tweet is None:
280 printNicely("-- None --")
281 elif tweet is Timeout:
282 printNicely("-- Timeout --")
283 elif tweet is HeartbeatTimeout:
284 printNicely("-- Heartbeat Timeout --")
285 elif tweet is Hangup:
286 printNicely("-- Hangup --")
287 elif tweet.get('text'):
288 draw(t=tweet)
289
290
291 def fly():
292 """
293 Main function
294 """
295 get_decorated_name()
296 p = Process(target=stream)
297 p.start()
298 g['stream_pid'] = p.pid
299 listen(sys.stdin)