684aa8ef6ad2173a5f1c2f56d38a28b688221c50
[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 + '] ['+ 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 line4 = ''
72
73 printNicely(line1)
74 printNicely(line2)
75 printNicely(line3)
76 printNicely('')
77
78
79 def parse_arguments():
80 """
81 Parse the arguments
82 """
83 parser = argparse.ArgumentParser(description=__doc__ or "")
84 parser.add_argument(
85 '-to',
86 '--timeout',
87 help='Timeout for the stream (seconds).')
88 parser.add_argument(
89 '-ht',
90 '--heartbeat-timeout',
91 help='Set heartbeat timeout.',
92 default=90)
93 parser.add_argument(
94 '-nb',
95 '--no-block',
96 action='store_true',
97 help='Set stream to non-blocking.')
98 parser.add_argument(
99 '-tt',
100 '--track-keywords',
101 help='Search the stream for specific text.')
102 return parser.parse_args()
103
104
105 def authen():
106 """
107 authenticate with Twitter OAuth
108 """
109 # When using rainbow stream you must authorize.
110 twitter_credential = os.environ.get(
111 'HOME',
112 os.environ.get(
113 'USERPROFILE',
114 '')) + os.sep + '.rainbow_oauth'
115 if not os.path.exists(twitter_credential):
116 oauth_dance("Rainbow Stream",
117 CONSUMER_KEY,
118 CONSUMER_SECRET,
119 twitter_credential)
120 oauth_token, oauth_token_secret = read_token_file(twitter_credential)
121 return OAuth(
122 oauth_token,
123 oauth_token_secret,
124 CONSUMER_KEY,
125 CONSUMER_SECRET)
126
127
128 def get_decorated_name():
129 """
130 Beginning of every line
131 """
132 t = Twitter(auth=authen())
133 name = '@' + t.statuses.user_timeline()[-1]['user']['screen_name']
134 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
135
136
137 def tweet():
138 """
139 Authen and tweet
140 """
141 t = Twitter(auth=authen())
142 t.statuses.update(status=g['stuff'])
143
144
145 def timeline():
146 """
147 Authen and get timeline
148 """
149 t = Twitter(auth=authen())
150 count = HOME_TWEET_NUM
151 if g['stuff'].isdigit():
152 count = g['stuff']
153 for tweet in reversed(t.statuses.home_timeline(count=count)):
154 draw(t=tweet)
155
156
157 def search():
158 """
159 Authen and search
160 """
161 t = Twitter(auth=authen())
162 rel = t.search.tweets(q='#' + g['stuff'])['statuses']
163 h, w = os.popen('stty size', 'r').read().split()
164
165 printNicely(grey('*' * int(w) + '\n'))
166 print('Newest', SEARCH_MAX_RECORD, 'tweet: \n')
167 for i in xrange(5):
168 draw(t=rel[i], keyword=g['stuff'].strip())
169 printNicely(grey('*' * int(w) + '\n'))
170
171
172 def friend():
173 """
174 List of friend (following)
175 """
176 t = Twitter(auth=authen())
177 g['friends'] = t.friends.ids()['ids']
178 for i in g['friends']:
179 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
180 user = cycle_color('@' + screen_name)
181 print(user, end=' ')
182 print('\n')
183
184
185 def follower():
186 """
187 List of follower
188 """
189 t = Twitter(auth=authen())
190 g['followers'] = t.followers.ids()['ids']
191 for i in g['followers']:
192 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
193 user = cycle_color('@' + screen_name)
194 print(user, end=' ')
195 print('\n')
196
197
198 def help():
199 """
200 Print help
201 """
202 usage = '''
203 Hi boss! I'm ready to serve you right now!
204 ----------------------------------------------------
205 "t" at the beginning will tweet immediately
206 "tl" at the beginning will tweet immediately
207 "s" and follow by any word will search and return 5 newest tweet
208 "fr" will list out your following people
209 "fl" will list out your followers
210 "h" or "help" will print this help once again
211 "c" will clear the terminal
212 "q" will exit
213 ----------------------------------------------------
214 Have fun and hang tight!
215 '''
216 printNicely(usage)
217 sys.stdout.write(g['decorated_name'])
218
219
220 def clear():
221 """
222 Exit all
223 """
224 os.system('clear')
225
226
227 def quit():
228 """
229 Exit all
230 """
231 os.kill(g['stream_pid'], signal.SIGKILL)
232 sys.exit()
233
234
235 def process(cmd):
236 """
237 Process switch by start of line
238 """
239 return {
240 't' : tweet,
241 'tl' : timeline,
242 's' : search,
243 'fr' : friend,
244 'fl' : follower,
245 'h' : help,
246 'help' : help,
247 'c' : clear,
248 'q' : quit,
249 }.get(cmd, lambda: sys.stdout.write(g['decorated_name']))
250
251
252 def listen(stdin):
253 """
254 Listen to user's input
255 """
256 for line in iter(stdin.readline, ''):
257 try:
258 cmd = line.split()[0]
259 except:
260 cmd = ''
261 # Save cmd to global variable and call process
262 g['stuff'] = ' '.join(line.split()[1:])
263 process(cmd)()
264 stdin.close()
265
266
267 def stream():
268 """
269 Track the stream
270 """
271 args = parse_arguments()
272
273 # The Logo
274 ascii_art()
275 # These arguments are optional:
276 stream_args = dict(
277 timeout=args.timeout,
278 block=not args.no_block,
279 heartbeat_timeout=args.heartbeat_timeout)
280
281 # Track keyword
282 query_args = dict()
283 if args.track_keywords:
284 query_args['track'] = args.track_keywords
285
286 # Get stream
287 stream = TwitterStream(
288 auth=authen(),
289 domain='userstream.twitter.com',
290 **stream_args)
291 tweet_iter = stream.user(**query_args)
292
293 # Iterate over the sample stream.
294 for tweet in tweet_iter:
295 if tweet is None:
296 printNicely("-- None --")
297 elif tweet is Timeout:
298 printNicely("-- Timeout --")
299 elif tweet is HeartbeatTimeout:
300 printNicely("-- Heartbeat Timeout --")
301 elif tweet is Hangup:
302 printNicely("-- Hangup --")
303 elif tweet.get('text'):
304 draw(t=tweet)
305
306
307 def fly():
308 """
309 Main function
310 """
311 get_decorated_name()
312 p = Process(target=stream)
313 p.start()
314 g['stream_pid'] = p.pid
315 listen(sys.stdin)