image viewer + input history + follow + unfollow
[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, imgflg = 0, 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 imgflg 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 '-img',
187 '--image',
188 help='Display all photo on terminal.')
189 return parser.parse_args()
190
191
192 def authen():
193 """
194 Authenticate with Twitter OAuth
195 """
196 # When using rainbow stream you must authorize.
197 twitter_credential = os.environ.get(
198 'HOME',
199 os.environ.get(
200 'USERPROFILE',
201 '')) + os.sep + '.rainbow_oauth'
202 if not os.path.exists(twitter_credential):
203 oauth_dance("Rainbow Stream",
204 CONSUMER_KEY,
205 CONSUMER_SECRET,
206 twitter_credential)
207 oauth_token, oauth_token_secret = read_token_file(twitter_credential)
208 return OAuth(
209 oauth_token,
210 oauth_token_secret,
211 CONSUMER_KEY,
212 CONSUMER_SECRET)
213
214
215 def get_decorated_name():
216 """
217 Beginning of every line
218 """
219 t = Twitter(auth=authen())
220 name = '@' + t.account.verify_credentials()['screen_name']
221 g['original_name'] = name[1:]
222 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
223
224
225 def switch():
226 """
227 Switch stream
228 """
229 try:
230 target = g['stuff'].split()[0]
231
232 # Filter and ignore
233 args = parse_arguments()
234 try:
235 if g['stuff'].split()[-1] == '-f':
236 only = raw_input('Only nicks: ')
237 ignore = raw_input('Ignore nicks: ')
238 args.filter = filter(None, only.split(','))
239 args.ignore = filter(None, ignore.split(','))
240 elif g['stuff'].split()[-1] == '-d':
241 args.filter = ONLY_LIST
242 args.ignore = IGNORE_LIST
243 except:
244 printNicely(red('Sorry, wrong format.'))
245 return
246
247 # Public stream
248 if target == 'public':
249 keyword = g['stuff'].split()[1]
250 if keyword[0] == '#':
251 keyword = keyword[1:]
252 # Kill old process
253 os.kill(g['stream_pid'], signal.SIGKILL)
254 args.track_keywords = keyword
255 # Start new process
256 p = Process(
257 target=stream,
258 args=(
259 PUBLIC_DOMAIN,
260 args))
261 p.start()
262 g['stream_pid'] = p.pid
263
264 # Personal stream
265 elif target == 'mine':
266 # Kill old process
267 os.kill(g['stream_pid'], signal.SIGKILL)
268 # Start new process
269 p = Process(
270 target=stream,
271 args=(
272 USER_DOMAIN,
273 args,
274 g['original_name']))
275 p.start()
276 g['stream_pid'] = p.pid
277 printNicely('')
278 printNicely(green('Stream switched.'))
279 if args.filter:
280 printNicely(cyan('Only: ' + str(args.filter)))
281 if args.ignore:
282 printNicely(red('Ignore: ' + str(args.ignore)))
283 printNicely('')
284 except:
285 printNicely(red('Sorry I can\'t understand.'))
286
287
288 def home():
289 """
290 Home
291 """
292 t = Twitter(auth=authen())
293 num = HOME_TWEET_NUM
294 if g['stuff'].isdigit():
295 num = g['stuff']
296 for tweet in reversed(t.statuses.home_timeline(count=num)):
297 draw(t=tweet, imgflg=g['image'])
298 printNicely('')
299
300
301 def view():
302 """
303 Friend view
304 """
305 t = Twitter(auth=authen())
306 user = g['stuff'].split()[0]
307 if user[0] == '@':
308 try:
309 num = int(g['stuff'].split()[1])
310 except:
311 num = HOME_TWEET_NUM
312 for tweet in reversed(t.statuses.user_timeline(count=num, screen_name=user[1:])):
313 draw(t=tweet, imgflg=g['image'])
314 printNicely('')
315 else:
316 printNicely(red('A name should begin with a \'@\''))
317
318
319 def tweet():
320 """
321 Tweet
322 """
323 t = Twitter(auth=authen())
324 t.statuses.update(status=g['stuff'])
325
326
327 def retweet():
328 """
329 ReTweet
330 """
331 t = Twitter(auth=authen())
332 try:
333 id = int(g['stuff'].split()[0])
334 tid = db.rainbow_query(id)[0].tweet_id
335 t.statuses.retweet(id=tid, include_entities=False, trim_user=True)
336 except:
337 printNicely(red('Sorry I can\'t retweet for you.'))
338
339
340 def favorite():
341 """
342 Favorite
343 """
344 t = Twitter(auth=authen())
345 try:
346 id = int(g['stuff'].split()[0])
347 tid = db.rainbow_query(id)[0].tweet_id
348 t.favorites.create(_id=tid, include_entities=False)
349 printNicely(green('Favorited.'))
350 draw(t.statuses.show(id=tid), imgflg=g['image'])
351 except:
352 printNicely(red('Omg some syntax is wrong.'))
353
354
355 def reply():
356 """
357 Reply
358 """
359 t = Twitter(auth=authen())
360 try:
361 id = int(g['stuff'].split()[0])
362 tid = db.rainbow_query(id)[0].tweet_id
363 user = t.statuses.show(id=tid)['user']['screen_name']
364 status = ' '.join(g['stuff'].split()[1:])
365 status = '@' + user + ' ' + status.decode('utf-8')
366 t.statuses.update(status=status, in_reply_to_status_id=tid)
367 except:
368 printNicely(red('Sorry I can\'t understand.'))
369
370
371 def delete():
372 """
373 Delete
374 """
375 t = Twitter(auth=authen())
376 try:
377 id = int(g['stuff'].split()[0])
378 tid = db.rainbow_query(id)[0].tweet_id
379 t.statuses.destroy(id=tid)
380 printNicely(green('Okay it\'s gone.'))
381 except:
382 printNicely(red('Sorry I can\'t delete this tweet for you.'))
383
384
385 def unfavorite():
386 """
387 Unfavorite
388 """
389 t = Twitter(auth=authen())
390 try:
391 id = int(g['stuff'].split()[0])
392 tid = db.rainbow_query(id)[0].tweet_id
393 t.favorites.destroy(_id=tid)
394 printNicely(green('Okay it\'s unfavorited.'))
395 draw(t.statuses.show(id=tid), imgflg=g['image'])
396 except:
397 printNicely(red('Sorry I can\'t unfavorite this tweet for you.'))
398
399
400 def search():
401 """
402 Search
403 """
404 t = Twitter(auth=authen())
405 try:
406 if g['stuff'][0] == '#':
407 rel = t.search.tweets(q=g['stuff'])['statuses']
408 if len(rel):
409 printNicely('Newest tweets:')
410 for i in reversed(xrange(SEARCH_MAX_RECORD)):
411 draw(t=rel[i],
412 imgflg=g['image'],
413 keyword=g['stuff'].strip()[1:])
414 printNicely('')
415 else:
416 printNicely(magenta('I\'m afraid there is no result'))
417 else:
418 printNicely(red('A keyword should be a hashtag (like \'#AKB48\')'))
419 except:
420 printNicely(red('Sorry I can\'t understand.'))
421
422
423 def show():
424 """
425 Show image
426 """
427 t = Twitter(auth=authen())
428 try:
429 target = g['stuff'].split()[0]
430 if target != 'image':
431 return
432 id = int(g['stuff'].split()[1])
433 tid = db.rainbow_query(id)[0].tweet_id
434 tweet = t.statuses.show(id=tid)
435 media = tweet['entities']['media']
436 for m in media:
437 res = requests.get(m['media_url'])
438 img = Image.open(StringIO(res.content))
439 img.show()
440 except:
441 printNicely(red('Sorry I can\'t show this image.'))
442
443
444 def follow():
445 """
446 Follow a user
447 """
448 t = Twitter(auth=authen())
449 screen_name = g['stuff'].split()[0]
450 if screen_name[0] == '@':
451 try :
452 t.friendships.create(screen_name=screen_name[1:],follow=True)
453 printNicely(green('You are following ' + screen_name + ' now!'))
454 except:
455 printNicely(red('Sorry can not follow at this time.'))
456 else:
457 printNicely(red('Sorry I can\'t understand.'))
458
459
460 def unfollow():
461 """
462 Unfollow a user
463 """
464 t = Twitter(auth=authen())
465 screen_name = g['stuff'].split()[0]
466 if screen_name[0] == '@':
467 try :
468 t.friendships.destroy(screen_name=screen_name[1:],include_entities=False)
469 printNicely(green('Unfollow ' + screen_name + ' success!'))
470 except:
471 printNicely(red('Sorry can not unfollow at this time.'))
472 else:
473 printNicely(red('Sorry I can\'t understand.'))
474
475
476 def help():
477 """
478 Help
479 """
480 s = ' ' * 2
481 h, w = os.popen('stty size', 'r').read().split()
482
483 usage = '\n'
484 usage += s + 'Hi boss! I\'m ready to serve you right now!\n'
485 usage += s + '-' * (int(w) - 4) + '\n'
486
487 usage += s + 'You are ' + yellow('already') + ' on your personal stream.\n'
488 usage += s * 2 + green('switch public #AKB') + \
489 ' will switch to public stream and follow "' + \
490 yellow('AKB') + '" keyword.\n'
491 usage += s * 2 + green('switch mine') + \
492 ' will switch to your personal stream.\n'
493 usage += s * 2 + green('switch mine -f ') + \
494 ' will prompt to enter the filter.\n'
495 usage += s * 3 + yellow('Only nicks') + \
496 ' filter will decide nicks will be INCLUDE ONLY.\n'
497 usage += s * 3 + yellow('Ignore nicks') + \
498 ' filter will decide nicks will be EXCLUDE.\n'
499 usage += s * 2 + green('switch mine -d') + \
500 ' will use the config\'s ONLY_LIST and IGNORE_LIST.\n'
501 usage += s * 3 + '(see ' + grey('rainbowstream/config.py') + ').\n'
502
503 usage += s + 'For more action: \n'
504 usage += s * 2 + green('home') + ' will show your timeline. ' + \
505 green('home 7') + ' will show 7 tweet.\n'
506 usage += s * 2 + green('view @mdo') + \
507 ' will show ' + yellow('@mdo') + '\'s home.\n'
508 usage += s * 2 + green('t oops ') + \
509 'will tweet "' + yellow('oops') + '" immediately.\n'
510 usage += s * 2 + \
511 green('rt 12 ') + ' will retweet to tweet with ' + \
512 yellow('[id=12]') + '.\n'
513 usage += s * 2 + \
514 green('fav 12 ') + ' will favorite the tweet with ' + \
515 yellow('[id=12]') + '.\n'
516 usage += s * 2 + green('rep 12 oops') + ' will reply "' + \
517 yellow('oops') + '" to tweet with ' + yellow('[id=12]') + '.\n'
518 usage += s * 2 + \
519 green('del 12 ') + ' will delete tweet with ' + \
520 yellow('[id=12]') + '.\n'
521 usage += s * 2 + \
522 green('ufav 12 ') + ' will unfavorite tweet with ' + \
523 yellow('[id=12]') + '.\n'
524 usage += s * 2 + green('s #AKB48') + ' will search for "' + \
525 yellow('AKB48') + '" and return 5 newest tweet.\n'
526 usage += s * 2 + green('show image 12') + ' will show image in ' + \
527 yellow('[id=12]') + 'in your OS\'s image viewer.\n'
528 usage += s * 2 + green('fl @dtvd88') + ' will follow ' + \
529 yellow('@dtvd88') + '.\n'
530 usage += s * 2 + green('ufl @dtvd88') + ' will unfollow ' + \
531 yellow('@dtvd88') + '.\n'
532 usage += s * 2 + green('h') + ' will show this help again.\n'
533 usage += s * 2 + green('c') + ' will clear the screen.\n'
534 usage += s * 2 + green('q') + ' will quit.\n'
535
536 usage += s + '-' * (int(w) - 4) + '\n'
537 usage += s + 'Have fun and hang tight!\n'
538 printNicely(usage)
539
540
541 def clear():
542 """
543 Clear screen
544 """
545 os.system('clear')
546
547
548 def quit():
549 """
550 Exit all
551 """
552 save_history()
553 os.system('rm -rf rainbow.db')
554 os.kill(g['stream_pid'], signal.SIGKILL)
555 sys.exit()
556
557
558 def reset():
559 """
560 Reset prefix of line
561 """
562 if g['reset']:
563 printNicely(magenta('Need tips ? Type "h" and hit Enter key!'))
564 g['reset'] = False
565
566
567 def process(cmd):
568 """
569 Process switch
570 """
571 return dict(zip(
572 cmdset,
573 [
574 switch,
575 home,
576 view,
577 tweet,
578 retweet,
579 favorite,
580 reply,
581 delete,
582 unfavorite,
583 search,
584 show,
585 follow,
586 unfollow,
587 help,
588 clear,
589 quit
590 ]
591 )).get(cmd, reset)
592
593
594 def listen():
595 """
596 Listen to user's input
597 """
598 d = dict(zip(
599 cmdset,
600 [
601 ['public #', 'mine'], # switch
602 [], # home
603 ['@'], # view
604 [], # tweet
605 [], # retweet
606 [], # favorite
607 [], # reply
608 [], # delete
609 [], # unfavorite
610 ['#'], # search
611 ['image'], # show image
612 [], # follow
613 [], # unfollow
614 [], # help
615 [], # clear
616 [], # quit
617 ]
618 ))
619 init_interactive_shell(d)
620 read_history()
621 reset()
622 while True:
623 if g['prefix']:
624 line = raw_input(g['decorated_name'])
625 else:
626 line = raw_input()
627 try:
628 cmd = line.split()[0]
629 except:
630 cmd = ''
631 # Save cmd to global variable and call process
632 g['stuff'] = ' '.join(line.split()[1:])
633 process(cmd)()
634 if cmd in ['switch', 't', 'rt', 'rep']:
635 g['prefix'] = False
636 else:
637 g['prefix'] = True
638
639
640 def stream(domain, args, name='Rainbow Stream'):
641 """
642 Track the stream
643 """
644
645 # The Logo
646 art_dict = {
647 USER_DOMAIN: name,
648 PUBLIC_DOMAIN: args.track_keywords,
649 SITE_DOMAIN: 'Site Stream',
650 }
651 ascii_art(art_dict[domain])
652
653 # These arguments are optional:
654 stream_args = dict(
655 timeout=args.timeout,
656 block=not args.no_block,
657 heartbeat_timeout=args.heartbeat_timeout)
658
659 # Track keyword
660 query_args = dict()
661 if args.track_keywords:
662 query_args['track'] = args.track_keywords
663
664 # Get stream
665 stream = TwitterStream(
666 auth=authen(),
667 domain=domain,
668 **stream_args)
669
670 if domain == USER_DOMAIN:
671 tweet_iter = stream.user(**query_args)
672 elif domain == SITE_DOMAIN:
673 tweet_iter = stream.site(**query_args)
674 else:
675 if args.track_keywords:
676 tweet_iter = stream.statuses.filter(**query_args)
677 else:
678 tweet_iter = stream.statuses.sample()
679
680 # Iterate over the stream.
681 for tweet in tweet_iter:
682 if tweet is None:
683 printNicely("-- None --")
684 elif tweet is Timeout:
685 printNicely("-- Timeout --")
686 elif tweet is HeartbeatTimeout:
687 printNicely("-- Heartbeat Timeout --")
688 elif tweet is Hangup:
689 printNicely("-- Hangup --")
690 elif tweet.get('text'):
691 draw(
692 t=tweet,
693 imgflg=args.image,
694 keyword=args.track_keywords,
695 fil=args.filter,
696 ig=args.ignore,
697 )
698
699
700 def fly():
701 """
702 Main function
703 """
704 # Spawn stream process
705 args = parse_arguments()
706 get_decorated_name()
707 p = Process(target=stream, args=(USER_DOMAIN, args, g['original_name']))
708 p.start()
709
710 # Start listen process
711 time.sleep(0.5)
712 g['reset'] = True
713 g['prefix'] = True
714 g['stream_pid'] = p.pid
715 g['image'] = args.image
716 listen()