clear will not add prefix
[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 reversed(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 g['prefix'] = False
320
321
322 def quit():
323 """
324 Exit all
325 """
326 os.system('rm -rf rainbow.db')
327 os.kill(g['stream_pid'], signal.SIGKILL)
328 sys.exit()
329
330
331 def reset():
332 """
333 Reset prefix of line
334 """
335 g['prefix'] = True
336
337
338 def process(cmd):
339 """
340 Process switch
341 """
342 return dict(zip(
343 cmdset,
344 [
345 home,
346 view,
347 tweet,
348 retweet,
349 reply,
350 delete,
351 search,
352 friend,
353 follower,
354 help,
355 clear,
356 quit
357 ]
358 )).get(cmd, reset)
359
360
361 def listen():
362 init_interactive_shell(cmdset)
363 first = True
364 while True:
365 if g['prefix'] and not first:
366 line = raw_input(g['decorated_name'])
367 else:
368 line = raw_input()
369 try:
370 cmd = line.split()[0]
371 except:
372 cmd = ''
373 # Save cmd to global variable and call process
374 g['stuff'] = ' '.join(line.split()[1:])
375 process(cmd)()
376 first = False
377
378
379 def stream():
380 """
381 Track the stream
382 """
383 args = parse_arguments()
384
385 # The Logo
386 ascii_art()
387 # These arguments are optional:
388 stream_args = dict(
389 timeout=args.timeout,
390 block=not args.no_block,
391 heartbeat_timeout=args.heartbeat_timeout)
392
393 # Track keyword
394 query_args = dict()
395 if args.track_keywords:
396 query_args['track'] = args.track_keywords
397
398 # Get stream
399 stream = TwitterStream(
400 auth=authen(),
401 domain=DOMAIN,
402 **stream_args)
403 tweet_iter = stream.user(**query_args)
404
405 # Iterate over the sample stream.
406 for tweet in tweet_iter:
407 if tweet is None:
408 printNicely("-- None --")
409 elif tweet is Timeout:
410 printNicely("-- Timeout --")
411 elif tweet is HeartbeatTimeout:
412 printNicely("-- Heartbeat Timeout --")
413 elif tweet is Hangup:
414 printNicely("-- Hangup --")
415 elif tweet.get('text'):
416 draw(t=tweet)
417
418
419 def fly():
420 """
421 Main function
422 """
423 get_decorated_name()
424 g['prefix'] = True
425 p = Process(target=stream)
426 p.start()
427 g['stream_pid'] = p.pid
428 listen()