e5476ba348db25e9dac25f7792aca4b622df9721
[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 .interactive import *
24 from .db import *
25
26 g = {}
27 db = RainbowDB()
28 cmdset = [
29 'switch',
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 ]
43
44
45 def draw(t, keyword=None, fil=[], ig=[]):
46 """
47 Draw the rainbow
48 """
49
50 # Retrieve tweet
51 tid = t['id']
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)
57 date = date - datetime.timedelta(seconds=time.timezone)
58 clock = date.strftime('%Y/%m/%d %H:%M:%S')
59
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
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
73 # Format info
74 user = cycle_color(name) + grey(' ' + screen_name + ' ')
75 meta = grey('[' + clock + '] [id=' + str(rid) + ']')
76 tweet = text.split()
77 # Highlight RT
78 tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
79 # Highlight screen_name
80 tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
81 # Highlight link
82 tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
83 # Highlight search keyword
84 if keyword:
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 )
91 tweet = ' '.join(tweet)
92
93 # Draw rainbow
94 line1 = u"{u:>{uw}}:".format(
95 u=user,
96 uw=len(user) + 2,
97 )
98 line2 = u"{c:>{cw}}".format(
99 c=meta,
100 cw=len(meta) + 2,
101 )
102 line3 = ' ' + tweet
103
104 printNicely('')
105 printNicely(line1)
106 printNicely(line2)
107 printNicely(line3)
108
109
110 def parse_arguments():
111 """
112 Parse the arguments
113 """
114 parser = argparse.ArgumentParser(description=__doc__ or "")
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.')
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.')
141 return parser.parse_args()
142
143
144 def authen():
145 """
146 Authenticate with Twitter OAuth
147 """
148 # When using rainbow stream you must authorize.
149 twitter_credential = os.environ.get(
150 'HOME',
151 os.environ.get(
152 'USERPROFILE',
153 '')) + os.sep + '.rainbow_oauth'
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)
160 return OAuth(
161 oauth_token,
162 oauth_token_secret,
163 CONSUMER_KEY,
164 CONSUMER_SECRET)
165
166
167 def get_decorated_name():
168 """
169 Beginning of every line
170 """
171 t = Twitter(auth=authen())
172 name = '@' + t.account.verify_credentials()['screen_name']
173 g['original_name'] = name[1:]
174 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
175
176
177 def switch():
178 """
179 Switch stream
180 """
181 try:
182 target = g['stuff'].split()[0]
183
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: ')
190 args.filter = filter(None,only.split(','))
191 args.ignore = filter(None,ignore.split(','))
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
199 # Public stream
200 if target == 'public':
201 keyword = g['stuff'].split()[1]
202 if keyword[0] == '#':
203 keyword = keyword[1:]
204 # Kill old process
205 os.kill(g['stream_pid'], signal.SIGKILL)
206 args.track_keywords = keyword
207 # Start new process
208 p = Process(
209 target=stream,
210 args=(
211 PUBLIC_DOMAIN,
212 args))
213 p.start()
214 g['stream_pid'] = p.pid
215
216 # Personal stream
217 elif target == 'mine':
218 # Kill old process
219 os.kill(g['stream_pid'], signal.SIGKILL)
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
229 printNicely('')
230 printNicely(green('Stream switched.'))
231 if args.filter:
232 printNicely(cyan('Only: ' + str(args.filter)))
233 if args.ignore:
234 printNicely(red('Ignore: ' + str(args.ignore)))
235 printNicely('')
236 except:
237 printNicely(red('Sorry I can\'t understand.'))
238
239
240 def home():
241 """
242 Home
243 """
244 t = Twitter(auth=authen())
245 num = HOME_TWEET_NUM
246 if g['stuff'].isdigit():
247 num = g['stuff']
248 for tweet in reversed(t.statuses.home_timeline(count=num)):
249 draw(t=tweet)
250 printNicely('')
251
252
253 def view():
254 """
255 Friend view
256 """
257 t = Twitter(auth=authen())
258 user = g['stuff'].split()[0]
259 if user[0] == '@':
260 try:
261 num = int(g['stuff'].split()[1])
262 except:
263 num = HOME_TWEET_NUM
264 for tweet in reversed(t.statuses.user_timeline(count=num, screen_name=user[1:])):
265 draw(t=tweet)
266 printNicely('')
267 else:
268 printNicely(red('A name should begin with a \'@\''))
269
270
271 def tweet():
272 """
273 Tweet
274 """
275 t = Twitter(auth=authen())
276 t.statuses.update(status=g['stuff'])
277
278
279 def 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
287 t.statuses.retweet(id=tid, include_entities=False, trim_user=True)
288 except:
289 printNicely(red('Sorry I can\'t retweet for you.'))
290
291
292 def reply():
293 """
294 Reply
295 """
296 t = Twitter(auth=authen())
297 try:
298 id = int(g['stuff'].split()[0])
299 tid = db.rainbow_query(id)[0].tweet_id
300 user = t.statuses.show(id=tid)['user']['screen_name']
301 status = ' '.join(g['stuff'].split()[1:])
302 status = '@' + user + ' ' + status.decode('utf-8')
303 t.statuses.update(status=status, in_reply_to_status_id=tid)
304 except:
305 printNicely(red('Sorry I can\'t understand.'))
306
307
308 def delete():
309 """
310 Delete
311 """
312 t = Twitter(auth=authen())
313 try:
314 id = int(g['stuff'].split()[0])
315 tid = db.rainbow_query(id)[0].tweet_id
316 t.statuses.destroy(id=tid)
317 printNicely(green('Okay it\'s gone.'))
318 except:
319 printNicely(red('Sorry I can\'t delete this tweet for you.'))
320
321
322 def search():
323 """
324 Search
325 """
326 t = Twitter(auth=authen())
327 try:
328 if g['stuff'][0] == '#':
329 rel = t.search.tweets(q=g['stuff'])['statuses']
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'))
337 else:
338 printNicely(red('A keyword should be a hashtag (like \'#AKB48\')'))
339 except:
340 printNicely(red('Sorry I can\'t understand.'))
341
342
343 def 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']:
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
354
355
356 def 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']:
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
367
368
369 def help():
370 """
371 Help
372 """
373 s = ' '*2
374 h,w = os.popen('stty size', 'r').read().split()
375
376 usage = '\n'
377 usage += s + 'Hi boss! I\'m ready to serve you right now!\n'
378 usage += s + '-'*(int(w)-4) + '\n'
379
380 usage += s + 'You are ' + yellow('already') +' on your personal stream.\n'
381 usage += s*2 + green('switch public #AKB') + ' will switch to public stream and follow "' + yellow('AKB') + '" keyword.\n'
382 usage += s*2 + green('switch mine') + ' will switch to your personal stream.\n'
383 usage += s*2 + green('switch mine -f ') + ' will prompt to enter the filter.\n'
384 usage += s*3 + yellow('Only nicks') + ' filter will decide nicks will be INCLUDE ONLY.\n'
385 usage += s*3 + yellow('Ignore nicks') + ' filter will decide nicks will be EXCLUDE.\n'
386 usage += s*2 + green('switch mine -d') + ' will use the config\'s ONLY_LIST and IGNORE_LIST.\n'
387 usage += s*3 + '(see ' + grey('rainbowstream/config.py') + ').\n'
388
389 usage += s + 'For more action: \n'
390 usage += s*2 + green('home') + ' will show your timeline. ' + green('home 7') + ' will show 7 tweet.\n'
391 usage += s*2 + green('view @mdo') + ' will show ' + yellow('@mdo') + '\'s home.\n'
392 usage += s*2 + green('t oops ') + 'will tweet "' + yellow('oops') + '" immediately.\n'
393 usage += s*2 + green('rt 12 ') + ' will retweet to tweet with '+ yellow('[id=12]') + '.\n'
394 usage += s*2 + green('rep 12 oops') + ' will reply "' + yellow('oops') + '" to tweet with ' + yellow('[id=12]') + '.\n'
395 usage += s*2 + green('del 12 ') + ' will delete tweet with ' + yellow('[id=12]') + '.\n'
396 usage += s*2 + green('s #AKB48') + ' will search for "'+yellow('AKB48') + '" and return 5 newest tweet.\n'
397 usage += s*2 + green('fr') + ' will list out your following people.\n'
398 usage += s*2 + green('fl') + ' will list out your follower.\n'
399 usage += s*2 + green('h') + ' will show this help again.\n'
400 usage += s*2 + green('c') + ' will clear the screen.\n'
401 usage += s*2 + green('q') + ' will quit.\n'
402
403 usage += s + '-'*(int(w)-4) + '\n'
404 usage += s + 'Have fun and hang tight!\n'
405 printNicely(usage)
406
407
408 def clear():
409 """
410 Clear screen
411 """
412 os.system('clear')
413
414
415 def quit():
416 """
417 Exit all
418 """
419 os.system('rm -rf rainbow.db')
420 os.kill(g['stream_pid'], signal.SIGKILL)
421 sys.exit()
422
423
424 def reset():
425 """
426 Reset prefix of line
427 """
428 if g['reset']:
429 printNicely(magenta('Need tips ? Type "h" and hit Enter key!'))
430 g['reset'] = False
431
432
433 def process(cmd):
434 """
435 Process switch
436 """
437 return dict(zip(
438 cmdset,
439 [
440 switch,
441 home,
442 view,
443 tweet,
444 retweet,
445 reply,
446 delete,
447 search,
448 friend,
449 follower,
450 help,
451 clear,
452 quit
453 ]
454 )).get(cmd, reset)
455
456
457 def listen():
458 """
459 Listen to user's input
460 """
461 d = dict(zip(
462 cmdset,
463 [
464 ['public #','mine'], # switch
465 [], # home
466 ['@'], # view
467 [], # tweet
468 [], # retweet
469 [], # reply
470 [], # delete
471 ['#'], # search
472 [], # friend
473 [], # follower
474 [], # help
475 [], # clear
476 [], # quit
477 ]
478 ))
479 init_interactive_shell(d)
480 reset()
481 while True:
482 if g['prefix']:
483 line = raw_input(g['decorated_name'])
484 else:
485 line = raw_input()
486 try:
487 cmd = line.split()[0]
488 except:
489 cmd = ''
490 # Save cmd to global variable and call process
491 g['stuff'] = ' '.join(line.split()[1:])
492 process(cmd)()
493 if cmd in ['switch','t','rt','rep']:
494 g['prefix'] = False
495 else:
496 g['prefix'] = True
497
498
499 def stream(domain, args, name='Rainbow Stream'):
500 """
501 Track the stream
502 """
503
504 # The Logo
505 art_dict = {
506 USER_DOMAIN: name,
507 PUBLIC_DOMAIN: args.track_keywords,
508 SITE_DOMAIN: 'Site Stream',
509 }
510 ascii_art(art_dict[domain])
511
512 # These arguments are optional:
513 stream_args = dict(
514 timeout=args.timeout,
515 block=not args.no_block,
516 heartbeat_timeout=args.heartbeat_timeout)
517
518 # Track keyword
519 query_args = dict()
520 if args.track_keywords:
521 query_args['track'] = args.track_keywords
522
523 # Get stream
524 stream = TwitterStream(
525 auth=authen(),
526 domain=domain,
527 **stream_args)
528
529 if domain == USER_DOMAIN:
530 tweet_iter = stream.user(**query_args)
531 elif domain == SITE_DOMAIN:
532 tweet_iter = stream.site(**query_args)
533 else:
534 if args.track_keywords:
535 tweet_iter = stream.statuses.filter(**query_args)
536 else:
537 tweet_iter = stream.statuses.sample()
538
539 # Iterate over the stream.
540 for tweet in tweet_iter:
541 if tweet is None:
542 printNicely("-- None --")
543 elif tweet is Timeout:
544 printNicely("-- Timeout --")
545 elif tweet is HeartbeatTimeout:
546 printNicely("-- Heartbeat Timeout --")
547 elif tweet is Hangup:
548 printNicely("-- Hangup --")
549 elif tweet.get('text'):
550 draw(t=tweet, keyword=args.track_keywords, fil=args.filter, ig=args.ignore)
551
552
553 def fly():
554 """
555 Main function
556 """
557 # Spawn stream process
558 args = parse_arguments()
559 get_decorated_name()
560 p = Process(target=stream, args=(USER_DOMAIN, args, g['original_name']))
561 p.start()
562
563 # Start listen process
564 time.sleep(0.5)
565 g['reset'] = True
566 g['prefix'] = True
567 g['stream_pid'] = p.pid
568 listen()
569