highlight
[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:]
42fde775 174 # Kill old process
175 os.kill(g['stream_pid'], signal.SIGKILL)
176 args = parse_arguments()
177 args.track_keywords = keyword
42fde775 178 # Start new process
179 p = Process(
180 target=stream,
181 args=(
182 PUBLIC_DOMAIN,
183 args))
184 p.start()
185 g['stream_pid'] = p.pid
186
187 # Personal stream
188 elif target == 'mine':
42fde775 189 # Kill old process
190 os.kill(g['stream_pid'], signal.SIGKILL)
191 args = parse_arguments()
42fde775 192 # Start new process
193 p = Process(
194 target=stream,
195 args=(
196 USER_DOMAIN,
197 args,
198 g['original_name']))
199 p.start()
200 g['stream_pid'] = p.pid
1551a7d3 201 printNicely(green('Stream switched.'))
42fde775 202 except:
203 printNicely(red('Sorry I can\'t understand.'))
204 g['prefix'] = False
205
206
7b674cef 207def home():
208 """
209 Home
210 """
211 t = Twitter(auth=authen())
94a5f62e 212 num = HOME_TWEET_NUM
7b674cef 213 if g['stuff'].isdigit():
94a5f62e 214 num = g['stuff']
215 for tweet in reversed(t.statuses.home_timeline(count=num)):
7b674cef 216 draw(t=tweet)
94a5f62e 217 printNicely('')
7b674cef 218
219
220def view():
221 """
222 Friend view
223 """
224 t = Twitter(auth=authen())
225 user = g['stuff'].split()[0]
b8fbcb70 226 if user[0] == '@':
227 try:
94a5f62e 228 num = int(g['stuff'].split()[1])
b8fbcb70 229 except:
94a5f62e 230 num = HOME_TWEET_NUM
231 for tweet in reversed(t.statuses.user_timeline(count=num, screen_name=user[1:])):
b8fbcb70 232 draw(t=tweet)
94a5f62e 233 printNicely('')
b8fbcb70 234 else:
c91f75f2 235 printNicely(red('A name should begin with a \'@\''))
7b674cef 236
237
f405a7d0 238def tweet():
54277114 239 """
7b674cef 240 Tweet
54277114
O
241 """
242 t = Twitter(auth=authen())
f405a7d0 243 t.statuses.update(status=g['stuff'])
94a5f62e 244 g['prefix'] = False
f405a7d0 245
b2b933a9 246
1ba4abfd
O
247def retweet():
248 """
249 ReTweet
250 """
251 t = Twitter(auth=authen())
252 try:
253 id = int(g['stuff'].split()[0])
254 tid = db.rainbow_query(id)[0].tweet_id
b2b933a9 255 t.statuses.retweet(id=tid, include_entities=False, trim_user=True)
1ba4abfd 256 except:
c91f75f2 257 printNicely(red('Sorry I can\'t retweet for you.'))
94a5f62e 258 g['prefix'] = False
1ba4abfd
O
259
260
7b674cef 261def reply():
829cc2d8 262 """
7b674cef 263 Reply
829cc2d8
O
264 """
265 t = Twitter(auth=authen())
7b674cef 266 try:
267 id = int(g['stuff'].split()[0])
18cab06a
O
268 tid = db.rainbow_query(id)[0].tweet_id
269 user = t.statuses.show(id=tid)['user']['screen_name']
7b674cef 270 status = ' '.join(g['stuff'].split()[1:])
271 status = '@' + user + ' ' + status.decode('utf-8')
18cab06a 272 t.statuses.update(status=status, in_reply_to_status_id=tid)
7b674cef 273 except:
c91f75f2 274 printNicely(red('Sorry I can\'t understand.'))
94a5f62e 275 g['prefix'] = False
7b674cef 276
277
278def delete():
279 """
280 Delete
281 """
282 t = Twitter(auth=authen())
283 try:
284 id = int(g['stuff'].split()[0])
18cab06a
O
285 tid = db.rainbow_query(id)[0].tweet_id
286 t.statuses.destroy(id=tid)
c91f75f2 287 printNicely(green('Okay it\'s gone.'))
7b674cef 288 except:
c91f75f2 289 printNicely(red('Sorry I can\'t delete this tweet for you.'))
829cc2d8
O
290
291
f405a7d0
O
292def search():
293 """
7b674cef 294 Search
f405a7d0
O
295 """
296 t = Twitter(auth=authen())
94a5f62e 297 try:
298 if g['stuff'][0] == '#':
299 rel = t.search.tweets(q=g['stuff'])['statuses']
c91f75f2 300 if len(rel):
301 printNicely('Newest tweets:')
302 for i in reversed(xrange(SEARCH_MAX_RECORD)):
303 draw(t=rel[i], keyword=g['stuff'].strip()[1:])
304 printNicely('')
305 else:
306 printNicely(magenta('I\'m afraid there is no result'))
94a5f62e 307 else:
c91f75f2 308 printNicely(red('A keyword should be a hashtag (like \'#AKB48\')'))
94a5f62e 309 except:
c91f75f2 310 printNicely(red('Sorry I can\'t understand.'))
b2b933a9 311
f405a7d0 312
843647ad
O
313def friend():
314 """
315 List of friend (following)
316 """
317 t = Twitter(auth=authen())
318 g['friends'] = t.friends.ids()['ids']
319 for i in g['friends']:
320 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
22be990e 321 user = cycle_color('@' + screen_name)
843647ad 322 print(user, end=' ')
42fde775 323 printNicely('')
843647ad
O
324
325
326def follower():
327 """
328 List of follower
329 """
330 t = Twitter(auth=authen())
331 g['followers'] = t.followers.ids()['ids']
332 for i in g['followers']:
333 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
22be990e 334 user = cycle_color('@' + screen_name)
843647ad 335 print(user, end=' ')
42fde775 336 printNicely('')
843647ad
O
337
338
f405a7d0
O
339def help():
340 """
7b674cef 341 Help
f405a7d0
O
342 """
343 usage = '''
344 Hi boss! I'm ready to serve you right now!
006f8cc4 345 -------------------------------------------------------------
d70182eb 346 You are already on your personal stream:
42fde775 347 "switch public #AKB" will switch to public stream and follow "AKB" keyword.
348 "switch mine" will switch back to your personal stream.
349 For more action:
350 "home" will show your timeline. "home 7" will show 7 tweet.
351 "view @bob" will show your friend @bob's home.
352 "t oops" will tweet "oops" immediately.
353 "rt 12345" will retweet to tweet with id "12345".
354 "rep 12345 oops" will reply "oops" to tweet with id "12345".
355 "del 12345" will delete tweet with id "12345".
356 "s #AKB48" will search for "AKB48" and return 5 newest tweet.
357 "fr" will list out your following people.
358 "fl" will list out your followers.
359 "h" will show this help again.
360 "c" will clear the terminal.
361 "q" will exit.
006f8cc4 362 -------------------------------------------------------------
843647ad 363 Have fun and hang tight!
f405a7d0
O
364 '''
365 printNicely(usage)
f405a7d0
O
366
367
843647ad 368def clear():
f405a7d0 369 """
7b674cef 370 Clear screen
f405a7d0 371 """
843647ad 372 os.system('clear')
7186f557 373 g['prefix'] = False
f405a7d0
O
374
375
843647ad 376def quit():
b8dda704
O
377 """
378 Exit all
379 """
8e633322 380 os.system('rm -rf rainbow.db')
843647ad
O
381 os.kill(g['stream_pid'], signal.SIGKILL)
382 sys.exit()
b8dda704
O
383
384
94a5f62e 385def reset():
f405a7d0 386 """
94a5f62e 387 Reset prefix of line
f405a7d0 388 """
c91f75f2 389 if g['reset']:
390 printNicely(green('Need tips ? Type "h" and hit Enter key!'))
94a5f62e 391 g['prefix'] = True
c91f75f2 392 g['reset'] = False
54277114
O
393
394
94a5f62e 395def process(cmd):
54277114 396 """
94a5f62e 397 Process switch
54277114 398 """
94a5f62e 399 return dict(zip(
400 cmdset,
b2b933a9 401 [
42fde775 402 switch,
b2b933a9 403 home,
404 view,
405 tweet,
406 retweet,
407 reply,
408 delete,
409 search,
410 friend,
411 follower,
412 help,
413 clear,
414 quit
415 ]
94a5f62e 416 )).get(cmd, reset)
417
418
419def listen():
42fde775 420 """
421 Listen to user's input
422 """
94a5f62e 423 init_interactive_shell(cmdset)
b2b933a9 424 first = True
425 while True:
94a5f62e 426 if g['prefix'] and not first:
427 line = raw_input(g['decorated_name'])
428 else:
429 line = raw_input()
843647ad
O
430 try:
431 cmd = line.split()[0]
432 except:
433 cmd = ''
f405a7d0 434 # Save cmd to global variable and call process
843647ad
O
435 g['stuff'] = ' '.join(line.split()[1:])
436 process(cmd)()
94a5f62e 437 first = False
54277114
O
438
439
42fde775 440def stream(domain, args, name='Rainbow Stream'):
54277114 441 """
f405a7d0 442 Track the stream
54277114 443 """
54277114 444 # The Logo
42fde775 445 art_dict = {
446 USER_DOMAIN: name,
447 PUBLIC_DOMAIN: args.track_keywords,
448 SITE_DOMAIN: 'Site Stream',
449 }
450 ascii_art(art_dict[domain])
91476ec3
O
451 # These arguments are optional:
452 stream_args = dict(
453 timeout=args.timeout,
454 block=not args.no_block,
455 heartbeat_timeout=args.heartbeat_timeout)
456
457 # Track keyword
458 query_args = dict()
459 if args.track_keywords:
460 query_args['track'] = args.track_keywords
461
462 # Get stream
2a6238f5 463 stream = TwitterStream(
22be990e 464 auth=authen(),
42fde775 465 domain=domain,
2a6238f5 466 **stream_args)
91476ec3 467
42fde775 468 if domain == USER_DOMAIN:
469 tweet_iter = stream.user(**query_args)
470 elif domain == SITE_DOMAIN:
471 tweet_iter = stream.site(**query_args)
472 else:
473 if args.track_keywords:
474 tweet_iter = stream.statuses.filter(**query_args)
475 else:
476 tweet_iter = stream.statuses.sample()
477
478 # Iterate over the stream.
91476ec3
O
479 for tweet in tweet_iter:
480 if tweet is None:
481 printNicely("-- None --")
482 elif tweet is Timeout:
483 printNicely("-- Timeout --")
484 elif tweet is HeartbeatTimeout:
485 printNicely("-- Heartbeat Timeout --")
486 elif tweet is Hangup:
487 printNicely("-- Hangup --")
488 elif tweet.get('text'):
1551a7d3 489 draw(t=tweet, keyword=args.track_keywords)
54277114
O
490
491
492def fly():
493 """
494 Main function
495 """
42fde775 496 # Spawn stream process
497 args = parse_arguments()
54277114 498 get_decorated_name()
42fde775 499 p = Process(target=stream, args=(USER_DOMAIN, args, g['original_name']))
500 p.start()
501
502 # Start listen process
94a5f62e 503 g['prefix'] = True
c91f75f2 504 g['reset'] = True
f405a7d0 505 g['stream_pid'] = p.pid
94a5f62e 506 listen()
fef2b940 507
508
509def multiplexer():
510 """
511 Multiplexer
512 """
513 os.system('screen -S rainbow fly')