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