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