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