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