update ss
[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
f405a7d0 8import os, os.path, argparse, sys, time, signal
91476ec3
O
9
10from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
54277114 11from twitter.api import *
91476ec3 12from twitter.oauth import OAuth, read_token_file
8c840a83 13from twitter.oauth_dance import oauth_dance
91476ec3
O
14from twitter.util import printNicely
15from dateutil import parser
91476ec3 16
2a6238f5
O
17from .colors import *
18from .config import *
19
f405a7d0 20g = {}
91476ec3 21
b8dda704 22def draw(t,keyword=None):
91476ec3
O
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
54277114 35 user = cycle_color(name) + grey(' ' + '@' + screen_name + ' ')
2a6238f5 36 clock = grey('[' + time + ']')
8c840a83 37 tweet = text.split()
b8dda704 38 # Highlight RT
2a6238f5 39 tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
b8dda704 40 # Highlight screen_name
2a6238f5 41 tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
b8dda704 42 # Highlight link
2a6238f5 43 tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
b8dda704 44 # Highlight search keyword
7a431249
O
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)
8c840a83 47 tweet = ' '.join(tweet)
91476ec3
O
48
49 # Draw rainbow
06773ffe 50 line1 = u"{u:>{uw}}:".format(
2a6238f5
O
51 u=user,
52 uw=len(user) + 2,
91476ec3 53 )
06773ffe 54 line2 = u"{c:>{cw}}".format(
2a6238f5
O
55 c=clock,
56 cw=len(clock) + 2,
06773ffe
O
57 )
58 line3 = ' ' + tweet
59 line4 = '\n'
91476ec3 60
f405a7d0
O
61 printNicely(line1)
62 printNicely(line2)
63 printNicely(line3)
64 printNicely(line4)
91476ec3
O
65
66
67def parse_arguments():
68 """
69 Parse the arguments
70 """
91476ec3 71 parser = argparse.ArgumentParser(description=__doc__ or "")
2a6238f5
O
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.')
91476ec3
O
90 return parser.parse_args()
91
92
54277114
O
93def authen():
94 """
95 authenticate with Twitter OAuth
96 """
8c840a83 97 # When using rainbow stream you must authorize.
2a6238f5
O
98 twitter_credential = os.environ.get(
99 'HOME',
100 os.environ.get(
101 'USERPROFILE',
102 '')) + os.sep + '.rainbow_oauth'
8c840a83
O
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)
54277114 109 return OAuth(
2a6238f5
O
110 oauth_token,
111 oauth_token_secret,
112 CONSUMER_KEY,
113 CONSUMER_SECRET)
91476ec3 114
54277114
O
115
116def 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']
f405a7d0 122 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
54277114 123
f405a7d0
O
124
125def tweet():
54277114
O
126 """
127 Authen and tweet
128 """
129 t = Twitter(auth=authen())
f405a7d0
O
130 t.statuses.update(status=g['stuff'])
131
132
133def 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'))
05105565 140 print('Newest',SEARCH_MAX_RECORD, 'tweet: \n')
f405a7d0 141 for i in xrange(5):
b8dda704 142 draw(t=rel[i],keyword=g['stuff'].strip())
f405a7d0
O
143 printNicely(grey('**************************************************************************************\n'))
144
145
843647ad
O
146def 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
159def 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
f405a7d0
O
172def help():
173 """
174 Print help
175 """
176 usage = '''
177 Hi boss! I'm ready to serve you right now!
178 ----------------------------------------------------
843647ad
O
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
b8dda704 184 "c" will clear the terminal
f405a7d0
O
185 "q" will exit
186 ----------------------------------------------------
843647ad 187 Have fun and hang tight!
f405a7d0
O
188 '''
189 printNicely(usage)
190 sys.stdout.write(g['decorated_name'])
191
192
843647ad 193def clear():
f405a7d0
O
194 """
195 Exit all
196 """
843647ad 197 os.system('clear')
f405a7d0
O
198
199
843647ad 200def quit():
b8dda704
O
201 """
202 Exit all
203 """
843647ad
O
204 os.kill(g['stream_pid'], signal.SIGKILL)
205 sys.exit()
b8dda704
O
206
207
843647ad 208def process(cmd):
f405a7d0
O
209 """
210 Process switch by start of line
211 """
212 return {
843647ad
O
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']))
54277114
O
221
222
223def listen(stdin):
224 """
225 Listen to user's input
226 """
227 for line in iter(stdin.readline, ''):
843647ad
O
228 try:
229 cmd = line.split()[0]
230 except:
231 cmd = ''
f405a7d0 232 # Save cmd to global variable and call process
843647ad
O
233 g['stuff'] = ' '.join(line.split()[1:])
234 process(cmd)()
54277114
O
235 stdin.close()
236
237
238def stream():
239 """
f405a7d0 240 Track the stream
54277114
O
241 """
242 args = parse_arguments()
243
244 # The Logo
245 ascii_art()
91476ec3
O
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
2a6238f5 258 stream = TwitterStream(
54277114
O
259 auth = authen(),
260 domain = 'userstream.twitter.com',
2a6238f5 261 **stream_args)
91476ec3
O
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'):
f405a7d0 275 draw(t=tweet)
54277114
O
276
277
278def fly():
279 """
280 Main function
281 """
282 get_decorated_name()
283 p = Process(target = stream)
284 p.start()
f405a7d0 285 g['stream_pid'] = p.pid
54277114
O
286 listen(sys.stdin)
287