1e50a0c2e9c87f9dbb8ee78f9038849decff4ce9
[rainbowstream.git] / rainbowstream / rainbow.py
1 """
2 Colorful user's timeline stream
3 """
4 from __future__ import print_function
5 from multiprocessing import Process
6 from dateutil import parser
7
8 import os
9 import os.path
10 import sys
11 import signal
12 import argparse
13 import time
14 import datetime
15
16 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
17 from twitter.api import *
18 from twitter.oauth import OAuth, read_token_file
19 from twitter.oauth_dance import oauth_dance
20 from twitter.util import printNicely
21
22 from .colors import *
23 from .config import *
24 from .interactive import *
25 from .db import *
26
27 g = {}
28 db = RainbowDB()
29 cmdset = [
30 'switch',
31 'home',
32 'view',
33 't',
34 'rt',
35 'rep',
36 'del',
37 's',
38 'fr',
39 'fl',
40 'h',
41 'c',
42 'q'
43 ]
44
45
46 def draw(t, keyword=None):
47 """
48 Draw the rainbow
49 """
50 # Retrieve tweet
51 tid = t['id']
52 text = t['text']
53 screen_name = t['user']['screen_name']
54 name = t['user']['name']
55 created_at = t['created_at']
56 date = parser.parse(created_at)
57 date = date - datetime.timedelta(seconds=time.timezone)
58 clock = date.strftime('%Y/%m/%d %H:%M:%S')
59
60 res = db.tweet_query(tid)
61 if not res:
62 db.store(tid)
63 res = db.tweet_query(tid)
64 rid = res[0].rainbow_id
65
66 # Format info
67 user = cycle_color(name) + grey(' ' + '@' + screen_name + ' ')
68 meta = grey('[' + clock + '] [id=' + str(rid) + ']')
69 tweet = text.split()
70 # Highlight RT
71 tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
72 # Highlight screen_name
73 tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
74 # Highlight link
75 tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
76 # Highlight search keyword
77 if keyword:
78 tweet = map(
79 lambda x: on_yellow(x) if
80 ''.join(c for c in x if c.isalnum()).lower() == keyword.lower()
81 else x,
82 tweet
83 )
84 tweet = ' '.join(tweet)
85
86 # Draw rainbow
87 line1 = u"{u:>{uw}}:".format(
88 u=user,
89 uw=len(user) + 2,
90 )
91 line2 = u"{c:>{cw}}".format(
92 c=meta,
93 cw=len(meta) + 2,
94 )
95 line3 = ' ' + tweet
96
97 printNicely('')
98 printNicely(line1)
99 printNicely(line2)
100 printNicely(line3)
101
102
103 def parse_arguments():
104 """
105 Parse the arguments
106 """
107 parser = argparse.ArgumentParser(description=__doc__ or "")
108 parser.add_argument(
109 '-to',
110 '--timeout',
111 help='Timeout for the stream (seconds).')
112 parser.add_argument(
113 '-ht',
114 '--heartbeat-timeout',
115 help='Set heartbeat timeout.',
116 default=90)
117 parser.add_argument(
118 '-nb',
119 '--no-block',
120 action='store_true',
121 help='Set stream to non-blocking.')
122 parser.add_argument(
123 '-tt',
124 '--track-keywords',
125 help='Search the stream for specific text.')
126 return parser.parse_args()
127
128
129 def authen():
130 """
131 Authenticate with Twitter OAuth
132 """
133 # When using rainbow stream you must authorize.
134 twitter_credential = os.environ.get(
135 'HOME',
136 os.environ.get(
137 'USERPROFILE',
138 '')) + os.sep + '.rainbow_oauth'
139 if not os.path.exists(twitter_credential):
140 oauth_dance("Rainbow Stream",
141 CONSUMER_KEY,
142 CONSUMER_SECRET,
143 twitter_credential)
144 oauth_token, oauth_token_secret = read_token_file(twitter_credential)
145 return OAuth(
146 oauth_token,
147 oauth_token_secret,
148 CONSUMER_KEY,
149 CONSUMER_SECRET)
150
151
152 def get_decorated_name():
153 """
154 Beginning of every line
155 """
156 t = Twitter(auth=authen())
157 name = '@' + t.account.verify_credentials()['screen_name']
158 g['original_name'] = name[1:]
159 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
160
161
162 def switch():
163 """
164 Switch stream
165 """
166 try:
167 target = g['stuff'].split()[0]
168
169 # Public stream
170 if target == 'public':
171 keyword = g['stuff'].split()[1]
172 if keyword[0] == '#':
173 keyword = keyword[1:]
174 # Kill old process
175 os.kill(g['stream_pid'], signal.SIGKILL)
176 args = parse_arguments()
177 args.track_keywords = keyword
178 # Start new process
179 p = Process(
180 target=stream,
181 args=(
182 PUBLIC_DOMAIN,
183 args))
184 p.start()
185 g['stream_pid'] = p.pid
186
187 # Personal stream
188 elif target == 'mine':
189 # Kill old process
190 os.kill(g['stream_pid'], signal.SIGKILL)
191 args = parse_arguments()
192 # Start new process
193 p = Process(
194 target=stream,
195 args=(
196 USER_DOMAIN,
197 args,
198 g['original_name']))
199 p.start()
200 g['stream_pid'] = p.pid
201 printNicely(green('Stream switched.'))
202 except:
203 printNicely(red('Sorry I can\'t understand.'))
204 g['prefix'] = False
205
206
207 def home():
208 """
209 Home
210 """
211 t = Twitter(auth=authen())
212 num = HOME_TWEET_NUM
213 if g['stuff'].isdigit():
214 num = g['stuff']
215 for tweet in reversed(t.statuses.home_timeline(count=num)):
216 draw(t=tweet)
217 printNicely('')
218
219
220 def view():
221 """
222 Friend view
223 """
224 t = Twitter(auth=authen())
225 user = g['stuff'].split()[0]
226 if user[0] == '@':
227 try:
228 num = int(g['stuff'].split()[1])
229 except:
230 num = HOME_TWEET_NUM
231 for tweet in reversed(t.statuses.user_timeline(count=num, screen_name=user[1:])):
232 draw(t=tweet)
233 printNicely('')
234 else:
235 printNicely(red('A name should begin with a \'@\''))
236
237
238 def tweet():
239 """
240 Tweet
241 """
242 t = Twitter(auth=authen())
243 t.statuses.update(status=g['stuff'])
244 g['prefix'] = False
245
246
247 def retweet():
248 """
249 ReTweet
250 """
251 t = Twitter(auth=authen())
252 try:
253 id = int(g['stuff'].split()[0])
254 tid = db.rainbow_query(id)[0].tweet_id
255 t.statuses.retweet(id=tid, include_entities=False, trim_user=True)
256 except:
257 printNicely(red('Sorry I can\'t retweet for you.'))
258 g['prefix'] = False
259
260
261 def reply():
262 """
263 Reply
264 """
265 t = Twitter(auth=authen())
266 try:
267 id = int(g['stuff'].split()[0])
268 tid = db.rainbow_query(id)[0].tweet_id
269 user = t.statuses.show(id=tid)['user']['screen_name']
270 status = ' '.join(g['stuff'].split()[1:])
271 status = '@' + user + ' ' + status.decode('utf-8')
272 t.statuses.update(status=status, in_reply_to_status_id=tid)
273 except:
274 printNicely(red('Sorry I can\'t understand.'))
275 g['prefix'] = False
276
277
278 def delete():
279 """
280 Delete
281 """
282 t = Twitter(auth=authen())
283 try:
284 id = int(g['stuff'].split()[0])
285 tid = db.rainbow_query(id)[0].tweet_id
286 t.statuses.destroy(id=tid)
287 printNicely(green('Okay it\'s gone.'))
288 except:
289 printNicely(red('Sorry I can\'t delete this tweet for you.'))
290
291
292 def search():
293 """
294 Search
295 """
296 t = Twitter(auth=authen())
297 try:
298 if g['stuff'][0] == '#':
299 rel = t.search.tweets(q=g['stuff'])['statuses']
300 if len(rel):
301 printNicely('Newest tweets:')
302 for i in reversed(xrange(SEARCH_MAX_RECORD)):
303 draw(t=rel[i], keyword=g['stuff'].strip()[1:])
304 printNicely('')
305 else:
306 printNicely(magenta('I\'m afraid there is no result'))
307 else:
308 printNicely(red('A keyword should be a hashtag (like \'#AKB48\')'))
309 except:
310 printNicely(red('Sorry I can\'t understand.'))
311
312
313 def friend():
314 """
315 List of friend (following)
316 """
317 t = Twitter(auth=authen())
318 g['friends'] = t.friends.ids()['ids']
319 for i in g['friends']:
320 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
321 user = cycle_color('@' + screen_name)
322 print(user, end=' ')
323 printNicely('')
324
325
326 def follower():
327 """
328 List of follower
329 """
330 t = Twitter(auth=authen())
331 g['followers'] = t.followers.ids()['ids']
332 for i in g['followers']:
333 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
334 user = cycle_color('@' + screen_name)
335 print(user, end=' ')
336 printNicely('')
337
338
339 def help():
340 """
341 Help
342 """
343 usage = '''
344 Hi boss! I'm ready to serve you right now!
345 -------------------------------------------------------------
346 You are already on your personal stream:
347 "switch public #AKB" will switch to public stream and follow "AKB" keyword.
348 "switch mine" will switch back to your personal stream.
349 For more action:
350 "home" will show your timeline. "home 7" will show 7 tweet.
351 "view @bob" will show your friend @bob's home.
352 "t oops" will tweet "oops" immediately.
353 "rt 12345" will retweet to tweet with id "12345".
354 "rep 12345 oops" will reply "oops" to tweet with id "12345".
355 "del 12345" will delete tweet with id "12345".
356 "s #AKB48" will search for "AKB48" and return 5 newest tweet.
357 "fr" will list out your following people.
358 "fl" will list out your followers.
359 "h" will show this help again.
360 "c" will clear the terminal.
361 "q" will exit.
362 -------------------------------------------------------------
363 Have fun and hang tight!
364 '''
365 printNicely(usage)
366
367
368 def clear():
369 """
370 Clear screen
371 """
372 os.system('clear')
373 g['prefix'] = False
374
375
376 def quit():
377 """
378 Exit all
379 """
380 os.system('rm -rf rainbow.db')
381 os.kill(g['stream_pid'], signal.SIGKILL)
382 sys.exit()
383
384
385 def reset():
386 """
387 Reset prefix of line
388 """
389 if g['reset']:
390 printNicely(green('Need tips ? Type "h" and hit Enter key!'))
391 g['prefix'] = True
392 g['reset'] = False
393
394
395 def process(cmd):
396 """
397 Process switch
398 """
399 return dict(zip(
400 cmdset,
401 [
402 switch,
403 home,
404 view,
405 tweet,
406 retweet,
407 reply,
408 delete,
409 search,
410 friend,
411 follower,
412 help,
413 clear,
414 quit
415 ]
416 )).get(cmd, reset)
417
418
419 def listen():
420 """
421 Listen to user's input
422 """
423 init_interactive_shell(cmdset)
424 first = True
425 while True:
426 if g['prefix'] and not first:
427 line = raw_input(g['decorated_name'])
428 else:
429 line = raw_input()
430 try:
431 cmd = line.split()[0]
432 except:
433 cmd = ''
434 # Save cmd to global variable and call process
435 g['stuff'] = ' '.join(line.split()[1:])
436 process(cmd)()
437 first = False
438
439
440 def stream(domain, args, name='Rainbow Stream'):
441 """
442 Track the stream
443 """
444 # The Logo
445 art_dict = {
446 USER_DOMAIN: name,
447 PUBLIC_DOMAIN: args.track_keywords,
448 SITE_DOMAIN: 'Site Stream',
449 }
450 ascii_art(art_dict[domain])
451 # These arguments are optional:
452 stream_args = dict(
453 timeout=args.timeout,
454 block=not args.no_block,
455 heartbeat_timeout=args.heartbeat_timeout)
456
457 # Track keyword
458 query_args = dict()
459 if args.track_keywords:
460 query_args['track'] = args.track_keywords
461
462 # Get stream
463 stream = TwitterStream(
464 auth=authen(),
465 domain=domain,
466 **stream_args)
467
468 if domain == USER_DOMAIN:
469 tweet_iter = stream.user(**query_args)
470 elif domain == SITE_DOMAIN:
471 tweet_iter = stream.site(**query_args)
472 else:
473 if args.track_keywords:
474 tweet_iter = stream.statuses.filter(**query_args)
475 else:
476 tweet_iter = stream.statuses.sample()
477
478 # Iterate over the stream.
479 for tweet in tweet_iter:
480 if tweet is None:
481 printNicely("-- None --")
482 elif tweet is Timeout:
483 printNicely("-- Timeout --")
484 elif tweet is HeartbeatTimeout:
485 printNicely("-- Heartbeat Timeout --")
486 elif tweet is Hangup:
487 printNicely("-- Hangup --")
488 elif tweet.get('text'):
489 draw(t=tweet, keyword=args.track_keywords)
490
491
492 def fly():
493 """
494 Main function
495 """
496 # Spawn stream process
497 args = parse_arguments()
498 get_decorated_name()
499 p = Process(target=stream, args=(USER_DOMAIN, args, g['original_name']))
500 p.start()
501
502 # Start listen process
503 g['prefix'] = True
504 g['reset'] = True
505 g['stream_pid'] = p.pid
506 listen()
507