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