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