control prefix
[rainbowstream.git] / rainbowstream / rainbow.py
1 """
2 Colorful user's timeline stream
3 """
4 from __future__ import print_function
5 from multiprocessing import Process
6 from dateutil import parser
7
8 import os
9 import os.path
10 import sys
11 import signal
12 import argparse
13 import time
14 import datetime
15
16 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
17 from twitter.api import *
18 from twitter.oauth import OAuth, read_token_file
19 from twitter.oauth_dance import oauth_dance
20 from twitter.util import printNicely
21
22 from .colors import *
23 from .config 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 'rep',
36 'del',
37 's',
38 'fr',
39 'fl',
40 'h',
41 'c',
42 'q'
43 ]
44
45
46 def draw(t, keyword=None, fil=[], ig=[]):
47 """
48 Draw the rainbow
49 """
50
51 # Retrieve tweet
52 tid = t['id']
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)
58 date = date - datetime.timedelta(seconds=time.timezone)
59 clock = date.strftime('%Y/%m/%d %H:%M:%S')
60
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
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
74 # Format info
75 user = cycle_color(name) + grey(' ' + screen_name + ' ')
76 meta = grey('[' + clock + '] [id=' + str(rid) + ']')
77 tweet = text.split()
78 # Highlight RT
79 tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
80 # Highlight screen_name
81 tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
82 # Highlight link
83 tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
84 # Highlight search keyword
85 if keyword:
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 )
92 tweet = ' '.join(tweet)
93
94 # Draw rainbow
95 line1 = u"{u:>{uw}}:".format(
96 u=user,
97 uw=len(user) + 2,
98 )
99 line2 = u"{c:>{cw}}".format(
100 c=meta,
101 cw=len(meta) + 2,
102 )
103 line3 = ' ' + tweet
104
105 printNicely('')
106 printNicely(line1)
107 printNicely(line2)
108 printNicely(line3)
109
110
111 def parse_arguments():
112 """
113 Parse the arguments
114 """
115 parser = argparse.ArgumentParser(description=__doc__ or "")
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.')
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.')
142 return parser.parse_args()
143
144
145 def authen():
146 """
147 Authenticate with Twitter OAuth
148 """
149 # When using rainbow stream you must authorize.
150 twitter_credential = os.environ.get(
151 'HOME',
152 os.environ.get(
153 'USERPROFILE',
154 '')) + os.sep + '.rainbow_oauth'
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)
161 return OAuth(
162 oauth_token,
163 oauth_token_secret,
164 CONSUMER_KEY,
165 CONSUMER_SECRET)
166
167
168 def get_decorated_name():
169 """
170 Beginning of every line
171 """
172 t = Twitter(auth=authen())
173 name = '@' + t.account.verify_credentials()['screen_name']
174 g['original_name'] = name[1:]
175 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
176
177
178 def switch():
179 """
180 Switch stream
181 """
182 try:
183 target = g['stuff'].split()[0]
184
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: ')
191 args.filter = only.split(',')[0]
192 args.ignore = ignore.split(',')[0]
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
200 # Public stream
201 if target == 'public':
202 keyword = g['stuff'].split()[1]
203 if keyword[0] == '#':
204 keyword = keyword[1:]
205 # Kill old process
206 os.kill(g['stream_pid'], signal.SIGKILL)
207 args.track_keywords = keyword
208 # Start new process
209 p = Process(
210 target=stream,
211 args=(
212 PUBLIC_DOMAIN,
213 args))
214 p.start()
215 g['stream_pid'] = p.pid
216
217 # Personal stream
218 elif target == 'mine':
219 # Kill old process
220 os.kill(g['stream_pid'], signal.SIGKILL)
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
230 printNicely('')
231 printNicely(green('Stream switched.'))
232 if args.filter:
233 printNicely(cyan('Only: ' + str(args.filter)))
234 if args.ignore:
235 printNicely(red('Ignore: ' + str(args.ignore)))
236 printNicely('')
237 except:
238 printNicely(red('Sorry I can\'t understand.'))
239
240
241 def home():
242 """
243 Home
244 """
245 t = Twitter(auth=authen())
246 num = HOME_TWEET_NUM
247 if g['stuff'].isdigit():
248 num = g['stuff']
249 for tweet in reversed(t.statuses.home_timeline(count=num)):
250 draw(t=tweet)
251 printNicely('')
252
253
254 def view():
255 """
256 Friend view
257 """
258 t = Twitter(auth=authen())
259 user = g['stuff'].split()[0]
260 if user[0] == '@':
261 try:
262 num = int(g['stuff'].split()[1])
263 except:
264 num = HOME_TWEET_NUM
265 for tweet in reversed(t.statuses.user_timeline(count=num, screen_name=user[1:])):
266 draw(t=tweet)
267 printNicely('')
268 else:
269 printNicely(red('A name should begin with a \'@\''))
270
271
272 def tweet():
273 """
274 Tweet
275 """
276 t = Twitter(auth=authen())
277 t.statuses.update(status=g['stuff'])
278
279
280 def 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
288 t.statuses.retweet(id=tid, include_entities=False, trim_user=True)
289 except:
290 printNicely(red('Sorry I can\'t retweet for you.'))
291
292
293 def reply():
294 """
295 Reply
296 """
297 t = Twitter(auth=authen())
298 try:
299 id = int(g['stuff'].split()[0])
300 tid = db.rainbow_query(id)[0].tweet_id
301 user = t.statuses.show(id=tid)['user']['screen_name']
302 status = ' '.join(g['stuff'].split()[1:])
303 status = '@' + user + ' ' + status.decode('utf-8')
304 t.statuses.update(status=status, in_reply_to_status_id=tid)
305 except:
306 printNicely(red('Sorry I can\'t understand.'))
307
308
309 def delete():
310 """
311 Delete
312 """
313 t = Twitter(auth=authen())
314 try:
315 id = int(g['stuff'].split()[0])
316 tid = db.rainbow_query(id)[0].tweet_id
317 t.statuses.destroy(id=tid)
318 printNicely(green('Okay it\'s gone.'))
319 except:
320 printNicely(red('Sorry I can\'t delete this tweet for you.'))
321
322
323 def search():
324 """
325 Search
326 """
327 t = Twitter(auth=authen())
328 try:
329 if g['stuff'][0] == '#':
330 rel = t.search.tweets(q=g['stuff'])['statuses']
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'))
338 else:
339 printNicely(red('A keyword should be a hashtag (like \'#AKB48\')'))
340 except:
341 printNicely(red('Sorry I can\'t understand.'))
342
343
344 def 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']
352 user = cycle_color('@' + screen_name)
353 print(user, end=' ')
354 printNicely('')
355
356
357 def 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']
365 user = cycle_color('@' + screen_name)
366 print(user, end=' ')
367 printNicely('')
368
369
370 def help():
371 """
372 Help
373 """
374 usage = '''
375 Hi boss! I'm ready to serve you right now!
376 -------------------------------------------------------------
377 You are already on your personal stream:
378 "switch public #AKB" will switch to public stream and follow "AKB" keyword.
379 "switch mine" will switch back to your personal stream.
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.
383 "switch mine -d" will use the config's ONLY_LIST and IGNORE_LIST
384 (see rainbowstream/config.py).
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.
398 -------------------------------------------------------------
399 Have fun and hang tight!
400 '''
401 printNicely(usage)
402
403
404 def clear():
405 """
406 Clear screen
407 """
408 os.system('clear')
409
410
411 def quit():
412 """
413 Exit all
414 """
415 os.system('rm -rf rainbow.db')
416 os.kill(g['stream_pid'], signal.SIGKILL)
417 sys.exit()
418
419
420 def reset():
421 """
422 Reset prefix of line
423 """
424 if g['reset']:
425 printNicely(green('Need tips ? Type "h" and hit Enter key!'))
426 g['reset'] = False
427
428
429 def process(cmd):
430 """
431 Process switch
432 """
433 return dict(zip(
434 cmdset,
435 [
436 switch,
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 ]
450 )).get(cmd, reset)
451
452
453 def listen():
454 """
455 Listen to user's input
456 """
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)
476 reset()
477 while True:
478 if g['prefix']:
479 line = raw_input(g['decorated_name'])
480 else:
481 line = raw_input()
482 try:
483 cmd = line.split()[0]
484 except:
485 cmd = ''
486 # Save cmd to global variable and call process
487 g['stuff'] = ' '.join(line.split()[1:])
488 process(cmd)()
489 if cmd in ['switch','t','rt','rep']:
490 g['prefix'] = False
491 else:
492 g['prefix'] = True
493
494
495 def stream(domain, args, name='Rainbow Stream'):
496 """
497 Track the stream
498 """
499
500 # The Logo
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])
507
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
520 stream = TwitterStream(
521 auth=authen(),
522 domain=domain,
523 **stream_args)
524
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.
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'):
546 draw(t=tweet, keyword=args.track_keywords, fil=args.filter, ig=args.ignore)
547
548
549 def fly():
550 """
551 Main function
552 """
553 # Spawn stream process
554 args = parse_arguments()
555 get_decorated_name()
556 p = Process(target=stream, args=(USER_DOMAIN, args, g['original_name']))
557 p.start()
558
559 # Start listen process
560 time.sleep(0.5)
561 g['reset'] = True
562 g['prefix'] = True
563 g['stream_pid'] = p.pid
564 listen()
565