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