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