9461d22679a274e70a61a35e571d0f91ecd6ee11
[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, os.path, argparse, sys, time, signal
9
10 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
11 from twitter.api import *
12 from twitter.oauth import OAuth, read_token_file
13 from twitter.oauth_dance import oauth_dance
14 from twitter.util import printNicely
15 from dateutil import parser
16
17 from .colors import *
18 from .config import *
19
20 g = {}
21
22 def draw(t,keyword=None):
23 """
24 Draw the rainbow
25 """
26 # Retrieve tweet
27 text = t['text']
28 screen_name = t['user']['screen_name']
29 name = t['user']['name']
30 created_at = t['created_at']
31 date = parser.parse(created_at)
32 time = date.strftime('%Y/%m/%d %H:%M:%S')
33
34 # Format info
35 user = cycle_color(name) + grey(' ' + '@' + screen_name + ' ')
36 clock = grey('[' + time + ']')
37 tweet = text.split()
38 # Highlight RT
39 tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
40 # Highlight screen_name
41 tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
42 # Highlight link
43 tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
44 # Highlight search keyword
45 if keyword:
46 tweet = map(lambda x: on_yellow(x) if ''.join(c for c in x if c.isalnum()).lower() == keyword.lower() else x, tweet)
47 tweet = ' '.join(tweet)
48
49 # Draw rainbow
50 line1 = u"{u:>{uw}}:".format(
51 u=user,
52 uw=len(user) + 2,
53 )
54 line2 = u"{c:>{cw}}".format(
55 c=clock,
56 cw=len(clock) + 2,
57 )
58 line3 = ' ' + tweet
59 line4 = '\n'
60
61 printNicely(line1)
62 printNicely(line2)
63 printNicely(line3)
64 printNicely(line4)
65
66
67 def parse_arguments():
68 """
69 Parse the arguments
70 """
71 parser = argparse.ArgumentParser(description=__doc__ or "")
72 parser.add_argument(
73 '-to',
74 '--timeout',
75 help='Timeout for the stream (seconds).')
76 parser.add_argument(
77 '-ht',
78 '--heartbeat-timeout',
79 help='Set heartbeat timeout.',
80 default=90)
81 parser.add_argument(
82 '-nb',
83 '--no-block',
84 action='store_true',
85 help='Set stream to non-blocking.')
86 parser.add_argument(
87 '-tt',
88 '--track-keywords',
89 help='Search the stream for specific text.')
90 return parser.parse_args()
91
92
93 def authen():
94 """
95 authenticate with Twitter OAuth
96 """
97 # When using rainbow stream you must authorize.
98 twitter_credential = os.environ.get(
99 'HOME',
100 os.environ.get(
101 'USERPROFILE',
102 '')) + os.sep + '.rainbow_oauth'
103 if not os.path.exists(twitter_credential):
104 oauth_dance("Rainbow Stream",
105 CONSUMER_KEY,
106 CONSUMER_SECRET,
107 twitter_credential)
108 oauth_token, oauth_token_secret = read_token_file(twitter_credential)
109 return OAuth(
110 oauth_token,
111 oauth_token_secret,
112 CONSUMER_KEY,
113 CONSUMER_SECRET)
114
115
116 def get_decorated_name():
117 """
118 Beginning of every line
119 """
120 t = Twitter(auth=authen())
121 name = '@' + t.statuses.user_timeline()[-1]['user']['screen_name']
122 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
123
124
125 def tweet():
126 """
127 Authen and tweet
128 """
129 t = Twitter(auth=authen())
130 t.statuses.update(status=g['stuff'])
131
132
133 def search():
134 """
135 Authen and search
136 """
137 t = Twitter(auth=authen())
138 rel = t.search.tweets(q='#' + g['stuff'])['statuses']
139 printNicely(grey('**************************************************************************************\n'))
140 print('Newest',SEARCH_MAX_RECORD, 'tweet: \n')
141 for i in xrange(5):
142 draw(t=rel[i],keyword=g['stuff'].strip())
143 printNicely(grey('**************************************************************************************\n'))
144
145
146 def friend():
147 """
148 List of friend (following)
149 """
150 t = Twitter(auth=authen())
151 g['friends'] = t.friends.ids()['ids']
152 for i in g['friends']:
153 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
154 user = cycle_color('@'+screen_name)
155 print(user, end=' ')
156 print('\n');
157
158
159 def follower():
160 """
161 List of follower
162 """
163 t = Twitter(auth=authen())
164 g['followers'] = t.followers.ids()['ids']
165 for i in g['followers']:
166 screen_name = t.users.lookup(user_id=i)[0]['screen_name']
167 user = cycle_color('@'+screen_name)
168 print(user, end=' ')
169 print('\n');
170
171
172 def help():
173 """
174 Print help
175 """
176 usage = '''
177 Hi boss! I'm ready to serve you right now!
178 ----------------------------------------------------
179 "tweet" at the beginning will tweet immediately
180 "s" and follow by any word will search and return 5 newest tweet
181 "fr" will list out your following people
182 "fl" will list out your followers
183 "h" will print this help once again
184 "c" will clear the terminal
185 "q" will exit
186 ----------------------------------------------------
187 Have fun and hang tight!
188 '''
189 printNicely(usage)
190 sys.stdout.write(g['decorated_name'])
191
192
193 def clear():
194 """
195 Exit all
196 """
197 os.system('clear')
198
199
200 def quit():
201 """
202 Exit all
203 """
204 os.kill(g['stream_pid'], signal.SIGKILL)
205 sys.exit()
206
207
208 def process(cmd):
209 """
210 Process switch by start of line
211 """
212 return {
213 'tweet' : tweet,
214 's' : search,
215 'fr' : friend,
216 'fl' : follower,
217 'h' : help,
218 'c' : clear,
219 'q' : quit,
220 }.get(cmd,lambda: sys.stdout.write(g['decorated_name']))
221
222
223 def listen(stdin):
224 """
225 Listen to user's input
226 """
227 for line in iter(stdin.readline, ''):
228 try:
229 cmd = line.split()[0]
230 except:
231 cmd = ''
232 # Save cmd to global variable and call process
233 g['stuff'] = ' '.join(line.split()[1:])
234 process(cmd)()
235 stdin.close()
236
237
238 def stream():
239 """
240 Track the stream
241 """
242 args = parse_arguments()
243
244 # The Logo
245 ascii_art()
246 # These arguments are optional:
247 stream_args = dict(
248 timeout=args.timeout,
249 block=not args.no_block,
250 heartbeat_timeout=args.heartbeat_timeout)
251
252 # Track keyword
253 query_args = dict()
254 if args.track_keywords:
255 query_args['track'] = args.track_keywords
256
257 # Get stream
258 stream = TwitterStream(
259 auth = authen(),
260 domain = 'userstream.twitter.com',
261 **stream_args)
262 tweet_iter = stream.user(**query_args)
263
264 # Iterate over the sample stream.
265 for tweet in tweet_iter:
266 if tweet is None:
267 printNicely("-- None --")
268 elif tweet is Timeout:
269 printNicely("-- Timeout --")
270 elif tweet is HeartbeatTimeout:
271 printNicely("-- Heartbeat Timeout --")
272 elif tweet is Hangup:
273 printNicely("-- Hangup --")
274 elif tweet.get('text'):
275 draw(t=tweet)
276
277
278 def fly():
279 """
280 Main function
281 """
282 get_decorated_name()
283 p = Process(target = stream)
284 p.start()
285 g['stream_pid'] = p.pid
286 listen(sys.stdin)
287