update
[rainbowstream.git] / rainbowstream / rainbow.py
CommitLineData
91476ec3
O
1"""
2Colorful user's timeline stream
3"""
4
5from __future__ import print_function
54277114 6from multiprocessing import Process
91476ec3 7
22be990e 8import os
9import os.path
10import argparse
11import sys
22be990e 12import signal
91476ec3
O
13
14from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
54277114 15from twitter.api import *
91476ec3 16from twitter.oauth import OAuth, read_token_file
8c840a83 17from twitter.oauth_dance import oauth_dance
91476ec3
O
18from twitter.util import printNicely
19from dateutil import parser
91476ec3 20
2a6238f5
O
21from .colors import *
22from .config import *
18cab06a 23from .db import *
2a6238f5 24
f405a7d0 25g = {}
18cab06a 26db = RainbowDB()
22be990e 27
28def draw(t, keyword=None):
91476ec3
O
29 """
30 Draw the rainbow
31 """
32 # Retrieve tweet
829cc2d8 33 tid = t['id']
91476ec3
O
34 text = t['text']
35 screen_name = t['user']['screen_name']
36 name = t['user']['name']
37 created_at = t['created_at']
38 date = parser.parse(created_at)
39 time = date.strftime('%Y/%m/%d %H:%M:%S')
40
18cab06a
O
41 res = db.tweet_query(tid)
42 if not res:
43 db.store(tid)
44 res = db.tweet_query(tid)
45 rid = res[0].rainbow_id
46
91476ec3 47 # Format info
54277114 48 user = cycle_color(name) + grey(' ' + '@' + screen_name + ' ')
18cab06a 49 meta = grey('[' + time + '] [id=' + str(rid) + ']')
8c840a83 50 tweet = text.split()
b8dda704 51 # Highlight RT
2a6238f5 52 tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
b8dda704 53 # Highlight screen_name
2a6238f5 54 tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
b8dda704 55 # Highlight link
2a6238f5 56 tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
b8dda704 57 # Highlight search keyword
7a431249 58 if keyword:
22be990e 59 tweet = map(
60 lambda x: on_yellow(x) if
61 ''.join(c for c in x if c.isalnum()).lower() == keyword.lower()
62 else x,
63 tweet
64 )
8c840a83 65 tweet = ' '.join(tweet)
91476ec3
O
66
67 # Draw rainbow
06773ffe 68 line1 = u"{u:>{uw}}:".format(
2a6238f5
O
69 u=user,
70 uw=len(user) + 2,
91476ec3 71 )
06773ffe 72 line2 = u"{c:>{cw}}".format(
829cc2d8
O
73 c=meta,
74 cw=len(meta) + 2,
06773ffe
O
75 )
76 line3 = ' ' + tweet
91476ec3 77
f405a7d0
O
78 printNicely(line1)
79 printNicely(line2)
80 printNicely(line3)
829cc2d8 81 printNicely('')
91476ec3
O
82
83
84def parse_arguments():
85 """
86 Parse the arguments
87 """
91476ec3 88 parser = argparse.ArgumentParser(description=__doc__ or "")
2a6238f5
O
89 parser.add_argument(
90 '-to',
91 '--timeout',
92 help='Timeout for the stream (seconds).')
93 parser.add_argument(
94 '-ht',
95 '--heartbeat-timeout',
96 help='Set heartbeat timeout.',
97 default=90)
98 parser.add_argument(
99 '-nb',
100 '--no-block',
101 action='store_true',
102 help='Set stream to non-blocking.')
103 parser.add_argument(
104 '-tt',
105 '--track-keywords',
106 help='Search the stream for specific text.')
91476ec3
O
107 return parser.parse_args()
108
109
54277114
O
110def authen():
111 """
7b674cef 112 Authenticate with Twitter OAuth
54277114 113 """
8c840a83 114 # When using rainbow stream you must authorize.
2a6238f5
O
115 twitter_credential = os.environ.get(
116 'HOME',
117 os.environ.get(
118 'USERPROFILE',
119 '')) + os.sep + '.rainbow_oauth'
8c840a83
O
120 if not os.path.exists(twitter_credential):
121 oauth_dance("Rainbow Stream",
122 CONSUMER_KEY,
123 CONSUMER_SECRET,
124 twitter_credential)
125 oauth_token, oauth_token_secret = read_token_file(twitter_credential)
54277114 126 return OAuth(
2a6238f5
O
127 oauth_token,
128 oauth_token_secret,
129 CONSUMER_KEY,
130 CONSUMER_SECRET)
91476ec3 131
54277114
O
132
133def get_decorated_name():
134 """
135 Beginning of every line
136 """
137 t = Twitter(auth=authen())
138 name = '@' + t.statuses.user_timeline()[-1]['user']['screen_name']
f405a7d0 139 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
54277114 140
f405a7d0 141
7b674cef 142def home():
143 """
144 Home
145 """
146 t = Twitter(auth=authen())
147 count = HOME_TWEET_NUM
148 if g['stuff'].isdigit():
149 count = g['stuff']
150 for tweet in reversed(t.statuses.home_timeline(count=count)):
151 draw(t=tweet)
152
153
154def view():
155 """
156 Friend view
157 """
158 t = Twitter(auth=authen())
159 user = g['stuff'].split()[0]
b8fbcb70 160 if user[0] == '@':
161 try:
162 count = int(g['stuff'].split()[1])
163 except:
164 count = HOME_TWEET_NUM
165 for tweet in reversed(t.statuses.user_timeline(count=count, screen_name=user[1:])):
166 draw(t=tweet)
167 else:
168 print(red('A name should begin with a \'@\''))
169 sys.stdout.write(g['decorated_name'])
7b674cef 170
171
f405a7d0 172def tweet():
54277114 173 """
7b674cef 174 Tweet
54277114
O
175 """
176 t = Twitter(auth=authen())
f405a7d0
O
177 t.statuses.update(status=g['stuff'])
178
179
1ba4abfd
O
180def retweet():
181 """
182 ReTweet
183 """
184 t = Twitter(auth=authen())
185 try:
186 id = int(g['stuff'].split()[0])
187 tid = db.rainbow_query(id)[0].tweet_id
188 t.statuses.retweet(id=tid,include_entities=False,trim_user=True)
189 except:
190 print(red('Sorry I can\'t retweet for you.'))
191 sys.stdout.write(g['decorated_name'])
192
193
7b674cef 194def reply():
829cc2d8 195 """
7b674cef 196 Reply
829cc2d8
O
197 """
198 t = Twitter(auth=authen())
7b674cef 199 try:
200 id = int(g['stuff'].split()[0])
18cab06a
O
201 tid = db.rainbow_query(id)[0].tweet_id
202 user = t.statuses.show(id=tid)['user']['screen_name']
7b674cef 203 status = ' '.join(g['stuff'].split()[1:])
204 status = '@' + user + ' ' + status.decode('utf-8')
18cab06a 205 t.statuses.update(status=status, in_reply_to_status_id=tid)
7b674cef 206 except:
207 print(red('Sorry I can\'t understand.'))
b8fbcb70 208 sys.stdout.write(g['decorated_name'])
7b674cef 209
210
211def delete():
212 """
213 Delete
214 """
215 t = Twitter(auth=authen())
216 try:
217 id = int(g['stuff'].split()[0])
18cab06a
O
218 tid = db.rainbow_query(id)[0].tweet_id
219 t.statuses.destroy(id=tid)
7b674cef 220 print(green('Okay it\'s gone.'))
221 except:
222 print(red('Sorry I can\'t delete this tweet for you.'))
18cab06a 223 sys.stdout.write(g['decorated_name'])
829cc2d8
O
224
225
f405a7d0
O
226def search():
227 """
7b674cef 228 Search
f405a7d0
O
229 """
230 t = Twitter(auth=authen())
2a911483 231 h, w = os.popen('stty size', 'r').read().split()
b8fbcb70 232 if g['stuff'][0] == '#':
233 rel = t.search.tweets(q=g['stuff'])['statuses']
234 printNicely(grey('*' * int(w) + '\n'))
235 print('Newest', SEARCH_MAX_RECORD, 'tweet: \n')
236 for i in xrange(5):
237 draw(t=rel[i], keyword=g['stuff'].strip())
238 printNicely(grey('*' * int(w) + '\n'))
239 else:
240 print(red('A keyword should be a hashtag (like \'#AKB48\')'))
241 sys.stdout.write(g['decorated_name'])
f405a7d0
O
242
243
843647ad
O
244def friend():
245 """
246 List of friend (following)
247 """
248 t = Twitter(auth=authen())
249 g['friends'] = t.friends.ids()['ids']
250 for i in g['friends']:
251 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
22be990e 252 user = cycle_color('@' + screen_name)
843647ad 253 print(user, end=' ')
22be990e 254 print('\n')
843647ad
O
255
256
257def follower():
258 """
259 List of follower
260 """
261 t = Twitter(auth=authen())
262 g['followers'] = t.followers.ids()['ids']
263 for i in g['followers']:
264 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
22be990e 265 user = cycle_color('@' + screen_name)
843647ad 266 print(user, end=' ')
22be990e 267 print('\n')
843647ad
O
268
269
f405a7d0
O
270def help():
271 """
7b674cef 272 Help
f405a7d0
O
273 """
274 usage = '''
275 Hi boss! I'm ready to serve you right now!
276 ----------------------------------------------------
c8f6a173 277 "home" will show your timeline. "home 7" will show 7 tweet.
b8fbcb70 278 "view @bob" will show your friend @bob's home.
7b674cef 279 "t oops" will tweet "oops" immediately.
1ba4abfd 280 "rt 12345" will retweet to tweet with id "12345".
a30eeccb 281 "rep 12345 oops" will reply "oops" to tweet with id "12345".
7b674cef 282 "del 12345" will delete tweet with id "12345".
b8fbcb70 283 "s #AKB48" will search for "AKB48" and return 5 newest tweet.
a30eeccb 284 "fr" will list out your following people.
285 "fl" will list out your followers.
c8f6a173 286 "h" will show this help again.
a30eeccb 287 "c" will clear the terminal.
288 "q" will exit.
f405a7d0 289 ----------------------------------------------------
843647ad 290 Have fun and hang tight!
f405a7d0
O
291 '''
292 printNicely(usage)
293 sys.stdout.write(g['decorated_name'])
294
295
843647ad 296def clear():
f405a7d0 297 """
7b674cef 298 Clear screen
f405a7d0 299 """
843647ad 300 os.system('clear')
f405a7d0
O
301
302
843647ad 303def quit():
b8dda704
O
304 """
305 Exit all
306 """
908d5012 307 db.truncate()
843647ad
O
308 os.kill(g['stream_pid'], signal.SIGKILL)
309 sys.exit()
b8dda704
O
310
311
843647ad 312def process(cmd):
f405a7d0 313 """
7b674cef 314 Process switch
f405a7d0
O
315 """
316 return {
1ba4abfd
O
317 'home' : home,
318 'view' : view,
319 't' : tweet,
320 'rt' : retweet,
321 'rep' : reply,
322 'del' : delete,
323 's' : search,
324 'fr' : friend,
325 'fl' : follower,
326 'h' : help,
327 'c' : clear,
328 'q' : quit,
22be990e 329 }.get(cmd, lambda: sys.stdout.write(g['decorated_name']))
54277114
O
330
331
332def listen(stdin):
333 """
334 Listen to user's input
335 """
336 for line in iter(stdin.readline, ''):
843647ad
O
337 try:
338 cmd = line.split()[0]
339 except:
340 cmd = ''
f405a7d0 341 # Save cmd to global variable and call process
843647ad
O
342 g['stuff'] = ' '.join(line.split()[1:])
343 process(cmd)()
54277114
O
344 stdin.close()
345
346
347def stream():
348 """
f405a7d0 349 Track the stream
54277114
O
350 """
351 args = parse_arguments()
352
353 # The Logo
354 ascii_art()
908d5012 355 g['stuff'] = '1'
356 home()
91476ec3
O
357 # These arguments are optional:
358 stream_args = dict(
359 timeout=args.timeout,
360 block=not args.no_block,
361 heartbeat_timeout=args.heartbeat_timeout)
362
363 # Track keyword
364 query_args = dict()
365 if args.track_keywords:
366 query_args['track'] = args.track_keywords
367
368 # Get stream
2a6238f5 369 stream = TwitterStream(
22be990e 370 auth=authen(),
c0440e50 371 domain=DOMAIN,
2a6238f5 372 **stream_args)
91476ec3
O
373 tweet_iter = stream.user(**query_args)
374
375 # Iterate over the sample stream.
376 for tweet in tweet_iter:
377 if tweet is None:
378 printNicely("-- None --")
379 elif tweet is Timeout:
380 printNicely("-- Timeout --")
381 elif tweet is HeartbeatTimeout:
382 printNicely("-- Heartbeat Timeout --")
383 elif tweet is Hangup:
384 printNicely("-- Hangup --")
385 elif tweet.get('text'):
f405a7d0 386 draw(t=tweet)
54277114
O
387
388
389def fly():
390 """
391 Main function
392 """
393 get_decorated_name()
18cab06a 394
22be990e 395 p = Process(target=stream)
54277114 396 p.start()
f405a7d0 397 g['stream_pid'] = p.pid
54277114 398 listen(sys.stdin)