add info for pypi
[rainbowstream.git] / rainbowstream / rainbow.py
1 """
2 Colorful user's timeline stream
3 """
4
5 from __future__ import print_function
6
7 import os, os.path, argparse
8
9 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
10 from twitter.oauth import OAuth, read_token_file
11 from twitter.oauth_dance import oauth_dance
12 from twitter.util import printNicely
13 from dateutil import parser
14
15 from colors import *
16 from config import *
17
18 def draw(t):
19 """
20 Draw the rainbow
21 """
22 # Retrieve tweet
23 text = t['text']
24 screen_name = t['user']['screen_name']
25 name = t['user']['name']
26 created_at = t['created_at']
27 date = parser.parse(created_at)
28 time = date.strftime('%Y/%m/%d %H:%M:%S')
29
30 # Format info
31 user = cycle_color(name + ' ' + '@' + screen_name + ' ')
32 clock = grey('['+ time + ']')
33 tweet = text.split()
34 tweet = map(lambda x: grey(x) if x=='RT' else x, tweet)
35 tweet = map(lambda x: cycle_color(x) if x[0]=='@' else x, tweet)
36 tweet = map(lambda x: cyan(x) if x[0:7]=='http://' else x, tweet)
37 tweet = ' '.join(tweet)
38
39 # Draw rainbow
40 line1 = u"{u:>{uw}}:".format(
41 u = user,
42 uw = len(user) + 2,
43 )
44 line2 = u"{c:>{cw}}".format(
45 c = clock,
46 cw = len(clock) + 2,
47 )
48 line3 = ' ' + tweet
49 line4 = '\n'
50
51 return line1, line2, line3, line4
52
53
54 def parse_arguments():
55 """
56 Parse the arguments
57 """
58
59 parser = argparse.ArgumentParser(description=__doc__ or "")
60
61 parser.add_argument('-to', '--timeout', help='Timeout for the stream (seconds).')
62 parser.add_argument('-ht', '--heartbeat-timeout', help='Set heartbeat timeout.', default=90)
63 parser.add_argument('-nb', '--no-block', action='store_true', help='Set stream to non-blocking.')
64 parser.add_argument('-tt', '--track-keywords', help='Search the stream for specific text.')
65 return parser.parse_args()
66
67
68 def stream():
69 args = parse_arguments()
70
71 # The Logo
72 ascii_art()
73
74 # When using rainbow stream you must authorize.
75 twitter_credential = os.environ.get('HOME', os.environ.get('USERPROFILE', '')) + os.sep + '.rainbow_oauth'
76 if not os.path.exists(twitter_credential):
77 oauth_dance("Rainbow Stream",
78 CONSUMER_KEY,
79 CONSUMER_SECRET,
80 twitter_credential)
81 oauth_token, oauth_token_secret = read_token_file(twitter_credential)
82 auth = OAuth(oauth_token, oauth_token_secret, CONSUMER_KEY, CONSUMER_SECRET)
83
84 # These arguments are optional:
85 stream_args = dict(
86 timeout=args.timeout,
87 block=not args.no_block,
88 heartbeat_timeout=args.heartbeat_timeout)
89
90 # Track keyword
91 query_args = dict()
92 if args.track_keywords:
93 query_args['track'] = args.track_keywords
94
95 # Get stream
96 stream = TwitterStream(auth=auth, domain='userstream.twitter.com', **stream_args)
97 tweet_iter = stream.user(**query_args)
98
99 # Iterate over the sample stream.
100 for tweet in tweet_iter:
101 if tweet is None:
102 printNicely("-- None --")
103 elif tweet is Timeout:
104 printNicely("-- Timeout --")
105 elif tweet is HeartbeatTimeout:
106 printNicely("-- Heartbeat Timeout --")
107 elif tweet is Hangup:
108 printNicely("-- Hangup --")
109 elif tweet.get('text'):
110 line1, line2, line3, line4 = draw(t = tweet)
111 printNicely(line1)
112 printNicely(line2)
113 printNicely(line3)
114 printNicely(line4)