9a30b47db5b8241f92dc69939c8f064b5c6810c2
[rainbowstream.git] / rainbowstream / rainbow.py
1 """
2 Colorful user's timeline stream
3 """
4
5 from __future__ import print_function
6 from multiprocessing import Process
7
8 import os
9 import os.path
10 import argparse
11 import sys
12 import signal
13
14 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
15 from twitter.api import *
16 from twitter.oauth import OAuth, read_token_file
17 from twitter.oauth_dance import oauth_dance
18 from twitter.util import printNicely
19 from dateutil import parser
20
21 from .colors import *
22 from .config import *
23 from .db import *
24
25 g = {}
26 db = RainbowDB()
27
28 def draw(t, keyword=None):
29 """
30 Draw the rainbow
31 """
32 # Retrieve tweet
33 tid = t['id']
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
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
47 # Format info
48 user = cycle_color(name) + grey(' ' + '@' + screen_name + ' ')
49 meta = grey('[' + time + '] [id=' + str(rid) + ']')
50 tweet = text.split()
51 # Highlight RT
52 tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
53 # Highlight screen_name
54 tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
55 # Highlight link
56 tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
57 # Highlight search keyword
58 if keyword:
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 )
65 tweet = ' '.join(tweet)
66
67 # Draw rainbow
68 line1 = u"{u:>{uw}}:".format(
69 u=user,
70 uw=len(user) + 2,
71 )
72 line2 = u"{c:>{cw}}".format(
73 c=meta,
74 cw=len(meta) + 2,
75 )
76 line3 = ' ' + tweet
77
78 printNicely(line1)
79 printNicely(line2)
80 printNicely(line3)
81 printNicely('')
82
83
84 def parse_arguments():
85 """
86 Parse the arguments
87 """
88 parser = argparse.ArgumentParser(description=__doc__ or "")
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.')
107 return parser.parse_args()
108
109
110 def authen():
111 """
112 Authenticate with Twitter OAuth
113 """
114 # When using rainbow stream you must authorize.
115 twitter_credential = os.environ.get(
116 'HOME',
117 os.environ.get(
118 'USERPROFILE',
119 '')) + os.sep + '.rainbow_oauth'
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)
126 return OAuth(
127 oauth_token,
128 oauth_token_secret,
129 CONSUMER_KEY,
130 CONSUMER_SECRET)
131
132
133 def 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']
139 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
140
141
142 def 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
154 def view():
155 """
156 Friend view
157 """
158 t = Twitter(auth=authen())
159 user = g['stuff'].split()[0]
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'])
170
171
172 def tweet():
173 """
174 Tweet
175 """
176 t = Twitter(auth=authen())
177 t.statuses.update(status=g['stuff'])
178
179
180 def 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
194 def reply():
195 """
196 Reply
197 """
198 t = Twitter(auth=authen())
199 try:
200 id = int(g['stuff'].split()[0])
201 tid = db.rainbow_query(id)[0].tweet_id
202 user = t.statuses.show(id=tid)['user']['screen_name']
203 status = ' '.join(g['stuff'].split()[1:])
204 status = '@' + user + ' ' + status.decode('utf-8')
205 t.statuses.update(status=status, in_reply_to_status_id=tid)
206 except:
207 print(red('Sorry I can\'t understand.'))
208 sys.stdout.write(g['decorated_name'])
209
210
211 def delete():
212 """
213 Delete
214 """
215 t = Twitter(auth=authen())
216 try:
217 id = int(g['stuff'].split()[0])
218 tid = db.rainbow_query(id)[0].tweet_id
219 t.statuses.destroy(id=tid)
220 print(green('Okay it\'s gone.'))
221 except:
222 print(red('Sorry I can\'t delete this tweet for you.'))
223 sys.stdout.write(g['decorated_name'])
224
225
226 def search():
227 """
228 Search
229 """
230 t = Twitter(auth=authen())
231 h, w = os.popen('stty size', 'r').read().split()
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'])
242
243
244 def 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']
252 user = cycle_color('@' + screen_name)
253 print(user, end=' ')
254 print('\n')
255
256
257 def 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']
265 user = cycle_color('@' + screen_name)
266 print(user, end=' ')
267 print('\n')
268
269
270 def help():
271 """
272 Help
273 """
274 usage = '''
275 Hi boss! I'm ready to serve you right now!
276 ----------------------------------------------------
277 "home" will show your timeline. "home 7" will print 7 tweet.
278 "view @bob" will show your friend @bob's home.
279 "t oops" will tweet "oops" immediately.
280 "rt 12345" will retweet to tweet with id "12345".
281 "rep 12345 oops" will reply "oops" to tweet with id "12345".
282 "del 12345" will delete tweet with id "12345".
283 "s #AKB48" will search for "AKB48" and return 5 newest tweet.
284 "fr" will list out your following people.
285 "fl" will list out your followers.
286 "h" will print this help once again.
287 "c" will clear the terminal.
288 "q" will exit.
289 ----------------------------------------------------
290 Have fun and hang tight!
291 '''
292 printNicely(usage)
293 sys.stdout.write(g['decorated_name'])
294
295
296 def clear():
297 """
298 Clear screen
299 """
300 os.system('clear')
301
302
303 def quit():
304 """
305 Exit all
306 """
307 os.kill(g['stream_pid'], signal.SIGKILL)
308 sys.exit()
309
310
311 def process(cmd):
312 """
313 Process switch
314 """
315 return {
316 'home' : home,
317 'view' : view,
318 't' : tweet,
319 'rt' : retweet,
320 'rep' : reply,
321 'del' : delete,
322 's' : search,
323 'fr' : friend,
324 'fl' : follower,
325 'h' : help,
326 'c' : clear,
327 'q' : quit,
328 }.get(cmd, lambda: sys.stdout.write(g['decorated_name']))
329
330
331 def listen(stdin):
332 """
333 Listen to user's input
334 """
335 for line in iter(stdin.readline, ''):
336 try:
337 cmd = line.split()[0]
338 except:
339 cmd = ''
340 # Save cmd to global variable and call process
341 g['stuff'] = ' '.join(line.split()[1:])
342 process(cmd)()
343 stdin.close()
344
345
346 def stream():
347 """
348 Track the stream
349 """
350 args = parse_arguments()
351
352 # The Logo
353 ascii_art()
354 # These arguments are optional:
355 stream_args = dict(
356 timeout=args.timeout,
357 block=not args.no_block,
358 heartbeat_timeout=args.heartbeat_timeout)
359
360 # Track keyword
361 query_args = dict()
362 if args.track_keywords:
363 query_args['track'] = args.track_keywords
364
365 # Get stream
366 stream = TwitterStream(
367 auth=authen(),
368 domain=DOMAIN,
369 **stream_args)
370 tweet_iter = stream.user(**query_args)
371
372 # Iterate over the sample stream.
373 for tweet in tweet_iter:
374 if tweet is None:
375 printNicely("-- None --")
376 elif tweet is Timeout:
377 printNicely("-- Timeout --")
378 elif tweet is HeartbeatTimeout:
379 printNicely("-- Heartbeat Timeout --")
380 elif tweet is Hangup:
381 printNicely("-- Hangup --")
382 elif tweet.get('text'):
383 draw(t=tweet)
384
385
386 def fly():
387 """
388 Main function
389 """
390 get_decorated_name()
391
392 p = Process(target=stream)
393 p.start()
394 g['stream_pid'] = p.pid
395 listen(sys.stdin)