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