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