84b89ffcf6a8f3988bef4c0c72257dbf98641297
[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 tweet = map(lambda x: on_yellow(x) if ''.join(c for c in x if c.isalnum()).lower() == keyword.lower() else x, tweet)
46 tweet = ' '.join(tweet)
47
48 # Draw rainbow
49 line1 = u"{u:>{uw}}:".format(
50 u=user,
51 uw=len(user) + 2,
52 )
53 line2 = u"{c:>{cw}}".format(
54 c=clock,
55 cw=len(clock) + 2,
56 )
57 line3 = ' ' + tweet
58 line4 = '\n'
59
60 printNicely(line1)
61 printNicely(line2)
62 printNicely(line3)
63 printNicely(line4)
64
65
66 def parse_arguments():
67 """
68 Parse the arguments
69 """
70
71 parser = argparse.ArgumentParser(description=__doc__ or "")
72
73 parser.add_argument(
74 '-to',
75 '--timeout',
76 help='Timeout for the stream (seconds).')
77 parser.add_argument(
78 '-ht',
79 '--heartbeat-timeout',
80 help='Set heartbeat timeout.',
81 default=90)
82 parser.add_argument(
83 '-nb',
84 '--no-block',
85 action='store_true',
86 help='Set stream to non-blocking.')
87 parser.add_argument(
88 '-tt',
89 '--track-keywords',
90 help='Search the stream for specific text.')
91 return parser.parse_args()
92
93
94 def authen():
95 """
96 authenticate with Twitter OAuth
97 """
98 # When using rainbow stream you must authorize.
99 twitter_credential = os.environ.get(
100 'HOME',
101 os.environ.get(
102 'USERPROFILE',
103 '')) + os.sep + '.rainbow_oauth'
104 if not os.path.exists(twitter_credential):
105 oauth_dance("Rainbow Stream",
106 CONSUMER_KEY,
107 CONSUMER_SECRET,
108 twitter_credential)
109 oauth_token, oauth_token_secret = read_token_file(twitter_credential)
110 return OAuth(
111 oauth_token,
112 oauth_token_secret,
113 CONSUMER_KEY,
114 CONSUMER_SECRET)
115
116
117 def get_decorated_name():
118 """
119 Beginning of every line
120 """
121 t = Twitter(auth=authen())
122 name = '@' + t.statuses.user_timeline()[-1]['user']['screen_name']
123 g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
124
125
126 def tweet():
127 """
128 Authen and tweet
129 """
130 t = Twitter(auth=authen())
131 t.statuses.update(status=g['stuff'])
132
133
134 def search():
135 """
136 Authen and search
137 """
138 t = Twitter(auth=authen())
139 rel = t.search.tweets(q='#' + g['stuff'])['statuses']
140 printNicely(grey('**************************************************************************************\n'))
141 print('Newest',SEARCH_MAX_RECORD, 'tweet: \n')
142 for i in xrange(5):
143 draw(t=rel[i],keyword=g['stuff'].strip())
144 printNicely(grey('**************************************************************************************\n'))
145
146
147 def help():
148 """
149 Print help
150 """
151 usage = '''
152 Hi boss! I'm ready to serve you right now!
153 ----------------------------------------------------
154 "!" at the beginning will start to tweet immediately
155 "/" at the beginning will search for 5 newest tweet
156 "?" or "h" will print this help once again
157 "c" will clear the terminal
158 "q" will exit
159 ----------------------------------------------------
160 Hvae fun and hang tight!
161 '''
162 printNicely(usage)
163 sys.stdout.write(g['decorated_name'])
164
165
166 def quit():
167 """
168 Exit all
169 """
170 os.kill(g['stream_pid'], signal.SIGKILL)
171 sys.exit()
172
173
174 def clear():
175 """
176 Exit all
177 """
178 os.system('clear')
179
180
181 def process(line):
182 """
183 Process switch by start of line
184 """
185 return {
186 '!' : tweet,
187 '/' : search,
188 '?' : help,
189 'h' : help,
190 'c' : clear,
191 'q' : quit,
192 }.get(line[0],lambda: sys.stdout.write(g['decorated_name']))
193
194
195 def listen(stdin):
196 """
197 Listen to user's input
198 """
199 for line in iter(stdin.readline, ''):
200 # Save cmd to global variable and call process
201 g['stuff'] = line[1:]
202 process(line)()
203 stdin.close()
204
205
206 def stream():
207 """
208 Track the stream
209 """
210 args = parse_arguments()
211
212 # The Logo
213 ascii_art()
214 # These arguments are optional:
215 stream_args = dict(
216 timeout=args.timeout,
217 block=not args.no_block,
218 heartbeat_timeout=args.heartbeat_timeout)
219
220 # Track keyword
221 query_args = dict()
222 if args.track_keywords:
223 query_args['track'] = args.track_keywords
224
225 # Get stream
226 stream = TwitterStream(
227 auth = authen(),
228 domain = 'userstream.twitter.com',
229 **stream_args)
230 tweet_iter = stream.user(**query_args)
231
232 # Iterate over the sample stream.
233 for tweet in tweet_iter:
234 if tweet is None:
235 printNicely("-- None --")
236 elif tweet is Timeout:
237 printNicely("-- Timeout --")
238 elif tweet is HeartbeatTimeout:
239 printNicely("-- Heartbeat Timeout --")
240 elif tweet is Hangup:
241 printNicely("-- Hangup --")
242 elif tweet.get('text'):
243 draw(t=tweet)
244
245
246 def fly():
247 """
248 Main function
249 """
250 get_decorated_name()
251 p = Process(target = stream)
252 p.start()
253 g['stream_pid'] = p.pid
254 listen(sys.stdin)
255