dc714e6f016db4f6c80008af60cdfd1e7a8d06cc
[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 print(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 print(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 print(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 print(green('Okay it\'s gone.'))
241 except:
242 print(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 print('Newest', SEARCH_MAX_RECORD, 'tweet:')
254 for i in xrange(SEARCH_MAX_RECORD):
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.'))
261
262
263 def 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']
271 user = cycle_color('@' + screen_name)
272 print(user, end=' ')
273 print('\n')
274
275
276 def 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']
284 user = cycle_color('@' + screen_name)
285 print(user, end=' ')
286 print('\n')
287
288
289 def help():
290 """
291 Help
292 """
293 usage = '''
294 Hi boss! I'm ready to serve you right now!
295 -------------------------------------------------------------
296 "home" will show your timeline. "home 7" will show 7 tweet.
297 "view @bob" will show your friend @bob's home.
298 "t oops" will tweet "oops" immediately.
299 "rt 12345" will retweet to tweet with id "12345".
300 "rep 12345 oops" will reply "oops" to tweet with id "12345".
301 "del 12345" will delete tweet with id "12345".
302 "s #AKB48" will search for "AKB48" and return 5 newest tweet.
303 "fr" will list out your following people.
304 "fl" will list out your followers.
305 "h" will show this help again.
306 "c" will clear the terminal.
307 "q" will exit.
308 -------------------------------------------------------------
309 Have fun and hang tight!
310 '''
311 printNicely(usage)
312
313
314 def clear():
315 """
316 Clear screen
317 """
318 os.system('clear')
319
320
321 def quit():
322 """
323 Exit all
324 """
325 os.system('rm -rf rainbow.db')
326 os.kill(g['stream_pid'], signal.SIGKILL)
327 sys.exit()
328
329
330 def reset():
331 """
332 Reset prefix of line
333 """
334 g['prefix'] = True
335
336
337 def process(cmd):
338 """
339 Process switch
340 """
341 return dict(zip(
342 cmdset,
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 ]
357 )).get(cmd, reset)
358
359
360 def listen():
361 init_interactive_shell(cmdset)
362 first = True
363 while True:
364 if g['prefix'] and not first:
365 line = raw_input(g['decorated_name'])
366 else:
367 line = raw_input()
368 try:
369 cmd = line.split()[0]
370 except:
371 cmd = ''
372 # Save cmd to global variable and call process
373 g['stuff'] = ' '.join(line.split()[1:])
374 process(cmd)()
375 first = False
376
377
378 def stream():
379 """
380 Track the stream
381 """
382 args = parse_arguments()
383
384 # The Logo
385 ascii_art()
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
398 stream = TwitterStream(
399 auth=authen(),
400 domain=DOMAIN,
401 **stream_args)
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'):
415 draw(t=tweet)
416
417
418 def fly():
419 """
420 Main function
421 """
422 get_decorated_name()
423 g['prefix'] = True
424 p = Process(target=stream)
425 p.start()
426 g['stream_pid'] = p.pid
427 listen()