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