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