add first tip
[rainbowstream.git] / rainbowstream / rainbow.py
1 """
2 Colorful user's timeline stream
3 """
4 from __future__ import print_function
5 from multiprocessing import Process
6 from dateutil import parser
7
8 import os
9 import os.path
10 import sys
11 import signal
12 import argparse
13 import time
14 import datetime
15
16 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
17 from twitter.api import *
18 from twitter.oauth import OAuth, read_token_file
19 from twitter.oauth_dance import oauth_dance
20 from twitter.util import printNicely
21
22 from .colors import *
23 from .config import *
24 from .interactive import *
25 from .db import *
26
27 g = {}
28 db = RainbowDB()
29 cmdset = [
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 ]
43
44
45 def draw(t, keyword=None):
46 """
47 Draw the rainbow
48 """
49 # Retrieve tweet
50 tid = t['id']
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)
56 date = date - datetime.timedelta(seconds=time.timezone)
57 clock = date.strftime('%Y/%m/%d %H:%M:%S')
58
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
65 # Format info
66 user = cycle_color(name) + grey(' ' + '@' + screen_name + ' ')
67 meta = grey('[' + clock + '] [id=' + str(rid) + ']')
68 tweet = text.split()
69 # Highlight RT
70 tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
71 # Highlight screen_name
72 tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
73 # Highlight link
74 tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
75 # Highlight search keyword
76 if keyword:
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 )
83 tweet = ' '.join(tweet)
84
85 # Draw rainbow
86 line1 = u"{u:>{uw}}:".format(
87 u=user,
88 uw=len(user) + 2,
89 )
90 line2 = u"{c:>{cw}}".format(
91 c=meta,
92 cw=len(meta) + 2,
93 )
94 line3 = ' ' + tweet
95
96 printNicely('')
97 printNicely(line1)
98 printNicely(line2)
99 printNicely(line3)
100
101
102 def parse_arguments():
103 """
104 Parse the arguments
105 """
106 parser = argparse.ArgumentParser(description=__doc__ or "")
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.')
125 return parser.parse_args()
126
127
128 def authen():
129 """
130 Authenticate with Twitter OAuth
131 """
132 # When using rainbow stream you must authorize.
133 twitter_credential = os.environ.get(
134 'HOME',
135 os.environ.get(
136 'USERPROFILE',
137 '')) + os.sep + '.rainbow_oauth'
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)
144 return OAuth(
145 oauth_token,
146 oauth_token_secret,
147 CONSUMER_KEY,
148 CONSUMER_SECRET)
149
150
151 def get_decorated_name():
152 """
153 Beginning of every line
154 """
155 t = Twitter(auth=authen())
156 name = '@' + t.account.verify_credentials()['screen_name']
157 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
158
159
160 def home():
161 """
162 Home
163 """
164 t = Twitter(auth=authen())
165 num = HOME_TWEET_NUM
166 if g['stuff'].isdigit():
167 num = g['stuff']
168 for tweet in reversed(t.statuses.home_timeline(count=num)):
169 draw(t=tweet)
170 printNicely('')
171
172
173 def view():
174 """
175 Friend view
176 """
177 t = Twitter(auth=authen())
178 user = g['stuff'].split()[0]
179 if user[0] == '@':
180 try:
181 num = int(g['stuff'].split()[1])
182 except:
183 num = HOME_TWEET_NUM
184 for tweet in reversed(t.statuses.user_timeline(count=num, screen_name=user[1:])):
185 draw(t=tweet)
186 printNicely('')
187 else:
188 printNicely(red('A name should begin with a \'@\''))
189
190
191 def tweet():
192 """
193 Tweet
194 """
195 t = Twitter(auth=authen())
196 t.statuses.update(status=g['stuff'])
197 g['prefix'] = False
198
199
200 def 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
208 t.statuses.retweet(id=tid, include_entities=False, trim_user=True)
209 except:
210 printNicely(red('Sorry I can\'t retweet for you.'))
211 g['prefix'] = False
212
213
214 def reply():
215 """
216 Reply
217 """
218 t = Twitter(auth=authen())
219 try:
220 id = int(g['stuff'].split()[0])
221 tid = db.rainbow_query(id)[0].tweet_id
222 user = t.statuses.show(id=tid)['user']['screen_name']
223 status = ' '.join(g['stuff'].split()[1:])
224 status = '@' + user + ' ' + status.decode('utf-8')
225 t.statuses.update(status=status, in_reply_to_status_id=tid)
226 except:
227 printNicely(red('Sorry I can\'t understand.'))
228 g['prefix'] = False
229
230
231 def delete():
232 """
233 Delete
234 """
235 t = Twitter(auth=authen())
236 try:
237 id = int(g['stuff'].split()[0])
238 tid = db.rainbow_query(id)[0].tweet_id
239 t.statuses.destroy(id=tid)
240 printNicely(green('Okay it\'s gone.'))
241 except:
242 printNicely(red('Sorry I can\'t delete this tweet for you.'))
243
244
245 def search():
246 """
247 Search
248 """
249 t = Twitter(auth=authen())
250 try:
251 if g['stuff'][0] == '#':
252 rel = t.search.tweets(q=g['stuff'])['statuses']
253 if len(rel):
254 printNicely('Newest tweets:')
255 for i in reversed(xrange(SEARCH_MAX_RECORD)):
256 draw(t=rel[i], keyword=g['stuff'].strip()[1:])
257 printNicely('')
258 else:
259 printNicely(magenta('I\'m afraid there is no result'))
260 else:
261 printNicely(red('A keyword should be a hashtag (like \'#AKB48\')'))
262 except:
263 printNicely(red('Sorry I can\'t understand.'))
264
265
266 def friend():
267 """
268 List of friend (following)
269 """
270 t = Twitter(auth=authen())
271 g['friends'] = t.friends.ids()['ids']
272 for i in g['friends']:
273 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
274 user = cycle_color('@' + screen_name)
275 print(user, end=' ')
276 printNicely('');
277
278
279 def follower():
280 """
281 List of follower
282 """
283 t = Twitter(auth=authen())
284 g['followers'] = t.followers.ids()['ids']
285 for i in g['followers']:
286 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
287 user = cycle_color('@' + screen_name)
288 print(user, end=' ')
289 printNicely('');
290
291
292 def help():
293 """
294 Help
295 """
296 usage = '''
297 Hi boss! I'm ready to serve you right now!
298 -------------------------------------------------------------
299 "home" will show your timeline. "home 7" will show 7 tweet.
300 "view @bob" will show your friend @bob's home.
301 "t oops" will tweet "oops" immediately.
302 "rt 12345" will retweet to tweet with id "12345".
303 "rep 12345 oops" will reply "oops" to tweet with id "12345".
304 "del 12345" will delete tweet with id "12345".
305 "s #AKB48" will search for "AKB48" and return 5 newest tweet.
306 "fr" will list out your following people.
307 "fl" will list out your followers.
308 "h" will show this help again.
309 "c" will clear the terminal.
310 "q" will exit.
311 -------------------------------------------------------------
312 Have fun and hang tight!
313 '''
314 printNicely(usage)
315
316
317 def clear():
318 """
319 Clear screen
320 """
321 os.system('clear')
322 g['prefix'] = False
323
324
325 def quit():
326 """
327 Exit all
328 """
329 os.system('rm -rf rainbow.db')
330 os.kill(g['stream_pid'], signal.SIGKILL)
331 sys.exit()
332
333
334 def reset():
335 """
336 Reset prefix of line
337 """
338 if g['reset']:
339 printNicely(green('Need tips ? Type "h" and hit Enter key!'))
340 g['prefix'] = True
341 g['reset'] = False
342
343
344 def process(cmd):
345 """
346 Process switch
347 """
348 return dict(zip(
349 cmdset,
350 [
351 home,
352 view,
353 tweet,
354 retweet,
355 reply,
356 delete,
357 search,
358 friend,
359 follower,
360 help,
361 clear,
362 quit
363 ]
364 )).get(cmd, reset)
365
366
367 def listen():
368 init_interactive_shell(cmdset)
369 first = True
370 while True:
371 if g['prefix'] and not first:
372 line = raw_input(g['decorated_name'])
373 else:
374 line = raw_input()
375 try:
376 cmd = line.split()[0]
377 except:
378 cmd = ''
379 # Save cmd to global variable and call process
380 g['stuff'] = ' '.join(line.split()[1:])
381 process(cmd)()
382 first = False
383
384
385 def stream():
386 """
387 Track the stream
388 """
389 args = parse_arguments()
390
391 # The Logo
392 ascii_art()
393 # These arguments are optional:
394 stream_args = dict(
395 timeout=args.timeout,
396 block=not args.no_block,
397 heartbeat_timeout=args.heartbeat_timeout)
398
399 # Track keyword
400 query_args = dict()
401 if args.track_keywords:
402 query_args['track'] = args.track_keywords
403
404 # Get stream
405 stream = TwitterStream(
406 auth=authen(),
407 domain=DOMAIN,
408 **stream_args)
409 tweet_iter = stream.user(**query_args)
410
411 # Iterate over the sample stream.
412 for tweet in tweet_iter:
413 if tweet is None:
414 printNicely("-- None --")
415 elif tweet is Timeout:
416 printNicely("-- Timeout --")
417 elif tweet is HeartbeatTimeout:
418 printNicely("-- Heartbeat Timeout --")
419 elif tweet is Hangup:
420 printNicely("-- Hangup --")
421 elif tweet.get('text'):
422 draw(t=tweet)
423
424
425 def fly():
426 """
427 Main function
428 """
429 get_decorated_name()
430 g['prefix'] = True
431 g['reset'] = True
432 p = Process(target=stream)
433 p.start()
434 g['stream_pid'] = p.pid
435 listen()