Merge branch 'master' of github.com:DTVD/rainbowstream
[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
15 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
16 from twitter.api import *
17 from twitter.oauth import OAuth, read_token_file
18 from twitter.oauth_dance import oauth_dance
19 from twitter.util import printNicely
20
21 from .colors import *
22 from .config import *
23 from .consumer import *
24 from .interactive import *
25 from .db import *
26
27 g = {}
28 db = RainbowDB()
29 cmdset = [
30 'switch',
31 'home',
32 'view',
33 't',
34 'rt',
35 'fav',
36 'rep',
37 'del',
38 'ufav',
39 's',
40 'fr',
41 'fl',
42 'h',
43 'c',
44 'q'
45 ]
46
47
48 def draw(t, keyword=None, fil=[], ig=[]):
49 """
50 Draw the rainbow
51 """
52
53 # Retrieve tweet
54 tid = t['id']
55 text = t['text']
56 screen_name = t['user']['screen_name']
57 name = t['user']['name']
58 created_at = t['created_at']
59 favorited = t['favorited']
60 date = parser.parse(created_at)
61 date = date - datetime.timedelta(seconds=time.timezone)
62 clock = date.strftime('%Y/%m/%d %H:%M:%S')
63
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
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
77 # Format info
78 user = cycle_color(name) + grey(' ' + screen_name + ' ')
79 meta = grey('[' + clock + '] [id=' + str(rid) + '] ')
80 if favorited:
81 meta = meta + green(u'\u2605')
82 tweet = text.split()
83 # Highlight RT
84 tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
85 # Highlight screen_name
86 tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
87 # Highlight link
88 tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
89 # Highlight search keyword
90 if keyword:
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 )
97 tweet = ' '.join(tweet)
98
99 # Draw rainbow
100 line1 = u"{u:>{uw}}:".format(
101 u=user,
102 uw=len(user) + 2,
103 )
104 line2 = u"{c:>{cw}}".format(
105 c=meta,
106 cw=len(meta) + 2,
107 )
108 line3 = ' ' + tweet
109
110 printNicely('')
111 printNicely(line1)
112 printNicely(line2)
113 printNicely(line3)
114
115
116 def parse_arguments():
117 """
118 Parse the arguments
119 """
120 parser = argparse.ArgumentParser(description=__doc__ or "")
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.')
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.')
147 return parser.parse_args()
148
149
150 def authen():
151 """
152 Authenticate with Twitter OAuth
153 """
154 # When using rainbow stream you must authorize.
155 twitter_credential = os.environ.get(
156 'HOME',
157 os.environ.get(
158 'USERPROFILE',
159 '')) + os.sep + '.rainbow_oauth'
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)
166 return OAuth(
167 oauth_token,
168 oauth_token_secret,
169 CONSUMER_KEY,
170 CONSUMER_SECRET)
171
172
173 def get_decorated_name():
174 """
175 Beginning of every line
176 """
177 t = Twitter(auth=authen())
178 name = '@' + t.account.verify_credentials()['screen_name']
179 g['original_name'] = name[1:]
180 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
181
182
183 def switch():
184 """
185 Switch stream
186 """
187 try:
188 target = g['stuff'].split()[0]
189
190 # Filter and ignore
191 args = parse_arguments()
192 try:
193 if g['stuff'].split()[-1] == '-f':
194 only = raw_input('Only nicks: ')
195 ignore = raw_input('Ignore nicks: ')
196 args.filter = filter(None, only.split(','))
197 args.ignore = filter(None, ignore.split(','))
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
205 # Public stream
206 if target == 'public':
207 keyword = g['stuff'].split()[1]
208 if keyword[0] == '#':
209 keyword = keyword[1:]
210 # Kill old process
211 os.kill(g['stream_pid'], signal.SIGKILL)
212 args.track_keywords = keyword
213 # Start new process
214 p = Process(
215 target=stream,
216 args=(
217 PUBLIC_DOMAIN,
218 args))
219 p.start()
220 g['stream_pid'] = p.pid
221
222 # Personal stream
223 elif target == 'mine':
224 # Kill old process
225 os.kill(g['stream_pid'], signal.SIGKILL)
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
235 printNicely('')
236 printNicely(green('Stream switched.'))
237 if args.filter:
238 printNicely(cyan('Only: ' + str(args.filter)))
239 if args.ignore:
240 printNicely(red('Ignore: ' + str(args.ignore)))
241 printNicely('')
242 except:
243 printNicely(red('Sorry I can\'t understand.'))
244
245
246 def home():
247 """
248 Home
249 """
250 t = Twitter(auth=authen())
251 num = HOME_TWEET_NUM
252 if g['stuff'].isdigit():
253 num = g['stuff']
254 for tweet in reversed(t.statuses.home_timeline(count=num)):
255 draw(t=tweet)
256 printNicely('')
257
258
259 def view():
260 """
261 Friend view
262 """
263 t = Twitter(auth=authen())
264 user = g['stuff'].split()[0]
265 if user[0] == '@':
266 try:
267 num = int(g['stuff'].split()[1])
268 except:
269 num = HOME_TWEET_NUM
270 for tweet in reversed(t.statuses.user_timeline(count=num, screen_name=user[1:])):
271 draw(t=tweet)
272 printNicely('')
273 else:
274 printNicely(red('A name should begin with a \'@\''))
275
276
277 def tweet():
278 """
279 Tweet
280 """
281 t = Twitter(auth=authen())
282 t.statuses.update(status=g['stuff'])
283
284
285 def 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
293 t.statuses.retweet(id=tid, include_entities=False, trim_user=True)
294 except:
295 printNicely(red('Sorry I can\'t retweet for you.'))
296
297
298 def 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
313 def reply():
314 """
315 Reply
316 """
317 t = Twitter(auth=authen())
318 try:
319 id = int(g['stuff'].split()[0])
320 tid = db.rainbow_query(id)[0].tweet_id
321 user = t.statuses.show(id=tid)['user']['screen_name']
322 status = ' '.join(g['stuff'].split()[1:])
323 status = '@' + user + ' ' + status.decode('utf-8')
324 t.statuses.update(status=status, in_reply_to_status_id=tid)
325 except:
326 printNicely(red('Sorry I can\'t understand.'))
327
328
329 def delete():
330 """
331 Delete
332 """
333 t = Twitter(auth=authen())
334 try:
335 id = int(g['stuff'].split()[0])
336 tid = db.rainbow_query(id)[0].tweet_id
337 t.statuses.destroy(id=tid)
338 printNicely(green('Okay it\'s gone.'))
339 except:
340 printNicely(red('Sorry I can\'t delete this tweet for you.'))
341
342
343 def 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
358 def search():
359 """
360 Search
361 """
362 t = Twitter(auth=authen())
363 try:
364 if g['stuff'][0] == '#':
365 rel = t.search.tweets(q=g['stuff'])['statuses']
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'))
373 else:
374 printNicely(red('A keyword should be a hashtag (like \'#AKB48\')'))
375 except:
376 printNicely(red('Sorry I can\'t understand.'))
377
378
379 def 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']:
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
390
391
392 def 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']:
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
403
404
405 def help():
406 """
407 Help
408 """
409 s = ' ' * 2
410 h, w = os.popen('stty size', 'r').read().split()
411
412 usage = '\n'
413 usage += s + 'Hi boss! I\'m ready to serve you right now!\n'
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'
431
432 usage += s + 'For more action: \n'
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'
462 usage += s + 'Have fun and hang tight!\n'
463 printNicely(usage)
464
465
466 def clear():
467 """
468 Clear screen
469 """
470 os.system('clear')
471
472
473 def quit():
474 """
475 Exit all
476 """
477 os.system('rm -rf rainbow.db')
478 os.kill(g['stream_pid'], signal.SIGKILL)
479 sys.exit()
480
481
482 def reset():
483 """
484 Reset prefix of line
485 """
486 if g['reset']:
487 printNicely(magenta('Need tips ? Type "h" and hit Enter key!'))
488 g['reset'] = False
489
490
491 def process(cmd):
492 """
493 Process switch
494 """
495 return dict(zip(
496 cmdset,
497 [
498 switch,
499 home,
500 view,
501 tweet,
502 retweet,
503 favorite,
504 reply,
505 delete,
506 unfavorite,
507 search,
508 friend,
509 follower,
510 help,
511 clear,
512 quit
513 ]
514 )).get(cmd, reset)
515
516
517 def listen():
518 """
519 Listen to user's input
520 """
521 d = dict(zip(
522 cmdset,
523 [
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
537 ]
538 ))
539 init_interactive_shell(d)
540 reset()
541 while True:
542 if g['prefix']:
543 line = raw_input(g['decorated_name'])
544 else:
545 line = raw_input()
546 try:
547 cmd = line.split()[0]
548 except:
549 cmd = ''
550 # Save cmd to global variable and call process
551 g['stuff'] = ' '.join(line.split()[1:])
552 process(cmd)()
553 if cmd in ['switch', 't', 'rt', 'rep']:
554 g['prefix'] = False
555 else:
556 g['prefix'] = True
557
558
559 def stream(domain, args, name='Rainbow Stream'):
560 """
561 Track the stream
562 """
563
564 # The Logo
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])
571
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
584 stream = TwitterStream(
585 auth=authen(),
586 domain=domain,
587 **stream_args)
588
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.
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'):
610 draw(
611 t=tweet,
612 keyword=args.track_keywords,
613 fil=args.filter,
614 ig=args.ignore)
615
616
617 def fly():
618 """
619 Main function
620 """
621 # Spawn stream process
622 args = parse_arguments()
623 get_decorated_name()
624 p = Process(target=stream, args=(USER_DOMAIN, args, g['original_name']))
625 p.start()
626
627 # Start listen process
628 time.sleep(0.5)
629 g['reset'] = True
630 g['prefix'] = True
631 g['stream_pid'] = p.pid
632 listen()