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