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