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