3a6c7e39229a6d5efffc668fcd5b6cbc4467dbb2
[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 time
13 import signal
14
15 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
16 from twitter.api import *
17 from twitter.oauth import OAuth, read_token_file
18 from twitter.oauth_dance import oauth_dance
19 from twitter.util import printNicely
20 from dateutil import parser
21
22 from .colors import *
23 from .config import *
24
25 g = {}
26
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 # Format info
42 user = cycle_color(name) + grey(' ' + '@' + screen_name + ' ')
43 meta = grey('[' + time + '] [id=' + str(tid) + ']')
44 tweet = text.split()
45 # Highlight RT
46 tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
47 # Highlight screen_name
48 tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
49 # Highlight link
50 tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
51 # Highlight search keyword
52 if keyword:
53 tweet = map(
54 lambda x: on_yellow(x) if
55 ''.join(c for c in x if c.isalnum()).lower() == keyword.lower()
56 else x,
57 tweet
58 )
59 tweet = ' '.join(tweet)
60
61 # Draw rainbow
62 line1 = u"{u:>{uw}}:".format(
63 u=user,
64 uw=len(user) + 2,
65 )
66 line2 = u"{c:>{cw}}".format(
67 c=meta,
68 cw=len(meta) + 2,
69 )
70 line3 = ' ' + tweet
71
72 printNicely(line1)
73 printNicely(line2)
74 printNicely(line3)
75 printNicely('')
76
77
78 def parse_arguments():
79 """
80 Parse the arguments
81 """
82 parser = argparse.ArgumentParser(description=__doc__ or "")
83 parser.add_argument(
84 '-to',
85 '--timeout',
86 help='Timeout for the stream (seconds).')
87 parser.add_argument(
88 '-ht',
89 '--heartbeat-timeout',
90 help='Set heartbeat timeout.',
91 default=90)
92 parser.add_argument(
93 '-nb',
94 '--no-block',
95 action='store_true',
96 help='Set stream to non-blocking.')
97 parser.add_argument(
98 '-tt',
99 '--track-keywords',
100 help='Search the stream for specific text.')
101 return parser.parse_args()
102
103
104 def authen():
105 """
106 Authenticate with Twitter OAuth
107 """
108 # When using rainbow stream you must authorize.
109 twitter_credential = os.environ.get(
110 'HOME',
111 os.environ.get(
112 'USERPROFILE',
113 '')) + os.sep + '.rainbow_oauth'
114 if not os.path.exists(twitter_credential):
115 oauth_dance("Rainbow Stream",
116 CONSUMER_KEY,
117 CONSUMER_SECRET,
118 twitter_credential)
119 oauth_token, oauth_token_secret = read_token_file(twitter_credential)
120 return OAuth(
121 oauth_token,
122 oauth_token_secret,
123 CONSUMER_KEY,
124 CONSUMER_SECRET)
125
126
127 def get_decorated_name():
128 """
129 Beginning of every line
130 """
131 t = Twitter(auth=authen())
132 name = '@' + t.statuses.user_timeline()[-1]['user']['screen_name']
133 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
134
135
136 def home():
137 """
138 Home
139 """
140 t = Twitter(auth=authen())
141 count = HOME_TWEET_NUM
142 if g['stuff'].isdigit():
143 count = g['stuff']
144 for tweet in reversed(t.statuses.home_timeline(count=count)):
145 draw(t=tweet)
146
147
148 def view():
149 """
150 Friend view
151 """
152 t = Twitter(auth=authen())
153 user = g['stuff'].split()[0]
154 if user[0] == '@':
155 try:
156 count = int(g['stuff'].split()[1])
157 except:
158 count = HOME_TWEET_NUM
159 for tweet in reversed(t.statuses.user_timeline(count=count, screen_name=user[1:])):
160 draw(t=tweet)
161 else:
162 print(red('A name should begin with a \'@\''))
163 sys.stdout.write(g['decorated_name'])
164
165
166 def tweet():
167 """
168 Tweet
169 """
170 t = Twitter(auth=authen())
171 t.statuses.update(status=g['stuff'])
172
173
174 def reply():
175 """
176 Reply
177 """
178 t = Twitter(auth=authen())
179 try:
180 id = int(g['stuff'].split()[0])
181 user = t.statuses.show(id=id)['user']['screen_name']
182 status = ' '.join(g['stuff'].split()[1:])
183 status = '@' + user + ' ' + status.decode('utf-8')
184 t.statuses.update(status=status, in_reply_to_status_id=id)
185 except:
186 print(red('Sorry I can\'t understand.'))
187 sys.stdout.write(g['decorated_name'])
188
189
190 def delete():
191 """
192 Delete
193 """
194 t = Twitter(auth=authen())
195 try:
196 id = int(g['stuff'].split()[0])
197 t.statuses.destroy(id=id)
198 print(green('Okay it\'s gone.'))
199 except:
200 print(red('Sorry I can\'t delete this tweet for you.'))
201 sys.stdout.write(g['decorated_name'])
202
203
204 def search():
205 """
206 Search
207 """
208 t = Twitter(auth=authen())
209 h, w = os.popen('stty size', 'r').read().split()
210 if g['stuff'][0] == '#':
211 rel = t.search.tweets(q=g['stuff'])['statuses']
212 printNicely(grey('*' * int(w) + '\n'))
213 print('Newest', SEARCH_MAX_RECORD, 'tweet: \n')
214 for i in xrange(5):
215 draw(t=rel[i], keyword=g['stuff'].strip())
216 printNicely(grey('*' * int(w) + '\n'))
217 else:
218 print(red('A keyword should be a hashtag (like \'#AKB48\')'))
219 sys.stdout.write(g['decorated_name'])
220
221
222 def friend():
223 """
224 List of friend (following)
225 """
226 t = Twitter(auth=authen())
227 g['friends'] = t.friends.ids()['ids']
228 for i in g['friends']:
229 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
230 user = cycle_color('@' + screen_name)
231 print(user, end=' ')
232 print('\n')
233
234
235 def follower():
236 """
237 List of follower
238 """
239 t = Twitter(auth=authen())
240 g['followers'] = t.followers.ids()['ids']
241 for i in g['followers']:
242 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
243 user = cycle_color('@' + screen_name)
244 print(user, end=' ')
245 print('\n')
246
247
248 def help():
249 """
250 Help
251 """
252 usage = '''
253 Hi boss! I'm ready to serve you right now!
254 ----------------------------------------------------
255 "home" will show your timeline. "home 7" will print 7 tweet.
256 "view @bob" will show your friend @bob's home.
257 "t oops" will tweet "oops" immediately.
258 "rep 12345 oops" will reply "oops" to tweet with id "12345".
259 "del 12345" will delete tweet with id "12345".
260 "s #AKB48" will search for "AKB48" and return 5 newest tweet.
261 "fr" will list out your following people.
262 "fl" will list out your followers.
263 "h" or "help" will print this help once again.
264 "c" will clear the terminal.
265 "q" will exit.
266 ----------------------------------------------------
267 Have fun and hang tight!
268 '''
269 printNicely(usage)
270 sys.stdout.write(g['decorated_name'])
271
272
273 def clear():
274 """
275 Clear screen
276 """
277 os.system('clear')
278
279
280 def quit():
281 """
282 Exit all
283 """
284 os.kill(g['stream_pid'], signal.SIGKILL)
285 sys.exit()
286
287
288 def process(cmd):
289 """
290 Process switch
291 """
292 return {
293 'home': home,
294 'view': view,
295 't': tweet,
296 'rep': reply,
297 'del': delete,
298 's': search,
299 'fr': friend,
300 'fl': follower,
301 'h': help,
302 'help': help,
303 'c': clear,
304 'q': quit,
305 }.get(cmd, lambda: sys.stdout.write(g['decorated_name']))
306
307
308 def listen(stdin):
309 """
310 Listen to user's input
311 """
312 for line in iter(stdin.readline, ''):
313 try:
314 cmd = line.split()[0]
315 except:
316 cmd = ''
317 # Save cmd to global variable and call process
318 g['stuff'] = ' '.join(line.split()[1:])
319 process(cmd)()
320 stdin.close()
321
322
323 def stream():
324 """
325 Track the stream
326 """
327 args = parse_arguments()
328
329 # The Logo
330 ascii_art()
331 # These arguments are optional:
332 stream_args = dict(
333 timeout=args.timeout,
334 block=not args.no_block,
335 heartbeat_timeout=args.heartbeat_timeout)
336
337 # Track keyword
338 query_args = dict()
339 if args.track_keywords:
340 query_args['track'] = args.track_keywords
341
342 # Get stream
343 stream = TwitterStream(
344 auth=authen(),
345 domain='userstream.twitter.com',
346 **stream_args)
347 tweet_iter = stream.user(**query_args)
348
349 # Iterate over the sample stream.
350 for tweet in tweet_iter:
351 if tweet is None:
352 printNicely("-- None --")
353 elif tweet is Timeout:
354 printNicely("-- Timeout --")
355 elif tweet is HeartbeatTimeout:
356 printNicely("-- Heartbeat Timeout --")
357 elif tweet is Hangup:
358 printNicely("-- Hangup --")
359 elif tweet.get('text'):
360 draw(t=tweet)
361
362
363 def fly():
364 """
365 Main function
366 """
367 get_decorated_name()
368 p = Process(target=stream)
369 p.start()
370 g['stream_pid'] = p.pid
371 listen(sys.stdin)