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
   9 import os.path
  10 import argparse
  11 import sys
  12 import time
  13 import signal
  14
  15 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
  16 from twitter.api import *
  17 from twitter.oauth import OAuth, read_token_file
  18 from twitter.oauth_dance import oauth_dance
  19 from twitter.util import printNicely
  20 from dateutil import parser
  21
  22 from .colors import *
  23 from .config import *
  24
  25 g = {}
  26
  27
  28 def draw(t, keyword=None):
  29     """
  30     Draw the rainbow
  31     """
  32     # Retrieve tweet
  33     text = t['text']
  34     screen_name = t['user']['screen_name']
  35     name = t['user']['name']
  36     created_at = t['created_at']
  37     date = parser.parse(created_at)
  38     time = date.strftime('%Y/%m/%d %H:%M:%S')
  39
  40     # Format info
  41     user = cycle_color(name) + grey(' ' + '@' + screen_name + ' ')
  42     clock = grey('[' + time + ']')
  43     tweet = text.split()
  44     # Highlight RT
  45     tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
  46     # Highlight screen_name
  47     tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
  48     # Highlight link
  49     tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
  50     # Highlight search keyword
  51     if keyword:
  52         tweet = map(
  53             lambda x: on_yellow(x) if
  54             ''.join(c for c in x if c.isalnum()).lower() == keyword.lower()
  55             else x,
  56             tweet
  57         )
  58     tweet = ' '.join(tweet)
  59
  60     # Draw rainbow
  61     line1 = u"{u:>{uw}}:".format(
  62         u=user,
  63         uw=len(user) + 2,
  64     )
  65     line2 = u"{c:>{cw}}".format(
  66         c=clock,
  67         cw=len(clock) + 2,
  68     )
  69     line3 = '  ' + tweet
  70     line4 = '\n'
  71
  72     printNicely(line1)
  73     printNicely(line2)
  74     printNicely(line3)
  75     printNicely(line4)
  76
  77
  78 def parse_arguments():
  79     """
  80     Parse the arguments
  81     """
  82     parser = argparse.ArgumentParser(description=__doc__ or "")
  83     parser.add_argument(
  84         '-to',
  85         '--timeout',
  86         help='Timeout for the stream (seconds).')
  87     parser.add_argument(
  88         '-ht',
  89         '--heartbeat-timeout',
  90         help='Set heartbeat timeout.',
  91         default=90)
  92     parser.add_argument(
  93         '-nb',
  94         '--no-block',
  95         action='store_true',
  96         help='Set stream to non-blocking.')
  97     parser.add_argument(
  98         '-tt',
  99         '--track-keywords',
 100         help='Search the stream for specific text.')
 101     return parser.parse_args()
 102
 103
 104 def authen():
 105     """
 106     authenticate with Twitter OAuth
 107     """
 108     # When using rainbow stream you must authorize.
 109     twitter_credential = os.environ.get(
 110         'HOME',
 111         os.environ.get(
 112             'USERPROFILE',
 113             '')) + os.sep + '.rainbow_oauth'
 114     if not os.path.exists(twitter_credential):
 115         oauth_dance("Rainbow Stream",
 116                     CONSUMER_KEY,
 117                     CONSUMER_SECRET,
 118                     twitter_credential)
 119     oauth_token, oauth_token_secret = read_token_file(twitter_credential)
 120     return OAuth(
 121         oauth_token,
 122         oauth_token_secret,
 123         CONSUMER_KEY,
 124         CONSUMER_SECRET)
 125
 126
 127 def get_decorated_name():
 128     """
 129     Beginning of every line
 130     """
 131     t = Twitter(auth=authen())
 132     name = '@' + t.statuses.user_timeline()[-1]['user']['screen_name']
 133     g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
 134
 135
 136 def tweet():
 137     """
 138     Authen and tweet
 139     """
 140     t = Twitter(auth=authen())
 141     t.statuses.update(status=g['stuff'])
 142
 143
 144 def search():
 145     """
 146     Authen and search
 147     """
 148     t = Twitter(auth=authen())
 149     rel = t.search.tweets(q='#' + g['stuff'])['statuses']
 150     h, w = os.popen('stty size', 'r').read().split()
 151
 152     printNicely(grey('*' * int(w) + '\n'))
 153     print('Newest', SEARCH_MAX_RECORD, 'tweet: \n')
 154     for i in xrange(5):
 155         draw(t=rel[i], keyword=g['stuff'].strip())
 156     printNicely(grey('*' * int(w) + '\n'))
 157
 158
 159 def friend():
 160     """
 161     List of friend (following)
 162     """
 163     t = Twitter(auth=authen())
 164     g['friends'] = t.friends.ids()['ids']
 165     for i in g['friends']:
 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
 172 def follower():
 173     """
 174     List of follower
 175     """
 176     t = Twitter(auth=authen())
 177     g['followers'] = t.followers.ids()['ids']
 178     for i in g['followers']:
 179         screen_name = t.users.lookup(user_id=i)[0]['screen_name']
 180         user = cycle_color('@' + screen_name)
 181         print(user, end=' ')
 182     print('\n')
 183
 184
 185 def help():
 186     """
 187     Print help
 188     """
 189     usage = '''
 190     Hi boss! I'm ready to serve you right now!
 191     ----------------------------------------------------
 192     "tweet" at the beginning will tweet immediately
 193     "s" and follow by any word will search and return 5 newest tweet
 194     "fr" will list out your following people
 195     "fl" will list out your followers
 196     "h" will print this help once again
 197     "c" will clear the terminal
 198     "q" will exit
 199     ----------------------------------------------------
 200     Have fun and hang tight!
 201     '''
 202     printNicely(usage)
 203     sys.stdout.write(g['decorated_name'])
 204
 205
 206 def clear():
 207     """
 208     Exit all
 209     """
 210     os.system('clear')
 211
 212
 213 def quit():
 214     """
 215     Exit all
 216     """
 217     os.kill(g['stream_pid'], signal.SIGKILL)
 218     sys.exit()
 219
 220
 221 def process(cmd):
 222     """
 223     Process switch by start of line
 224     """
 225     return {
 226         'tweet' : tweet,
 227         's'     : search,
 228         'fr'    : friend,
 229         'fl'    : follower,
 230         'h'     : help,
 231         'c'     : clear,
 232         'q'     : quit,
 233     }.get(cmd, lambda: sys.stdout.write(g['decorated_name']))
 234
 235
 236 def listen(stdin):
 237     """
 238     Listen to user's input
 239     """
 240     for line in iter(stdin.readline, ''):
 241         try:
 242             cmd = line.split()[0]
 243         except:
 244             cmd = ''
 245         # Save cmd to global variable and call process
 246         g['stuff'] = ' '.join(line.split()[1:])
 247         process(cmd)()
 248     stdin.close()
 249
 250
 251 def stream():
 252     """
 253     Track the stream
 254     """
 255     args = parse_arguments()
 256
 257     # The Logo
 258     ascii_art()
 259     # These arguments are optional:
 260     stream_args = dict(
 261         timeout=args.timeout,
 262         block=not args.no_block,
 263         heartbeat_timeout=args.heartbeat_timeout)
 264
 265     # Track keyword
 266     query_args = dict()
 267     if args.track_keywords:
 268         query_args['track'] = args.track_keywords
 269
 270     # Get stream
 271     stream = TwitterStream(
 272         auth=authen(),
 273         domain='userstream.twitter.com',
 274         **stream_args)
 275     tweet_iter = stream.user(**query_args)
 276
 277     # Iterate over the sample stream.
 278     for tweet in tweet_iter:
 279         if tweet is None:
 280             printNicely("-- None --")
 281         elif tweet is Timeout:
 282             printNicely("-- Timeout --")
 283         elif tweet is HeartbeatTimeout:
 284             printNicely("-- Heartbeat Timeout --")
 285         elif tweet is Hangup:
 286             printNicely("-- Hangup --")
 287         elif tweet.get('text'):
 288             draw(t=tweet)
 289
 290
 291 def fly():
 292     """
 293     Main function
 294     """
 295     get_decorated_name()
 296     p = Process(target=stream)
 297     p.start()
 298     g['stream_pid'] = p.pid
 299     listen(sys.stdin)