rainbowstream/rainbow.py

   1 """
   2 Colorful user's timeline stream
   3 """
   4
   5 from __future__ import print_function
   6
   7 import os
   8 import os.path
   9 import argparse
  10
  11 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
  12 from twitter.oauth import OAuth, read_token_file
  13 from twitter.oauth_dance import oauth_dance
  14 from twitter.util import printNicely
  15 from dateutil import parser
  16
  17 from .colors import *
  18 from .config import *
  19
  20
  21 def 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
  34     user = cycle_color(name + ' ' + '@' + screen_name + ' ')
  35     clock = grey('[' + time + ']')
  36     tweet = text.split()
  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)
  40     tweet = ' '.join(tweet)
  41
  42     # Draw rainbow
  43     line1 = u"{u:>{uw}}:".format(
  44         u=user,
  45         uw=len(user) + 2,
  46     )
  47     line2 = u"{c:>{cw}}".format(
  48         c=clock,
  49         cw=len(clock) + 2,
  50     )
  51     line3 = '  ' + tweet
  52     line4 = '\n'
  53
  54     return line1, line2, line3, line4
  55
  56
  57 def parse_arguments():
  58     """
  59     Parse the arguments
  60     """
  61
  62     parser = argparse.ArgumentParser(description=__doc__ or "")
  63
  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.')
  82     return parser.parse_args()
  83
  84
  85 def stream():
  86     args = parse_arguments()
  87
  88     # The Logo
  89     ascii_art()
  90
  91     # When using rainbow stream you must authorize.
  92     twitter_credential = os.environ.get(
  93         'HOME',
  94         os.environ.get(
  95             'USERPROFILE',
  96             '')) + os.sep + '.rainbow_oauth'
  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)
 103     auth = OAuth(
 104         oauth_token,
 105         oauth_token_secret,
 106         CONSUMER_KEY,
 107         CONSUMER_SECRET)
 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
 121     stream = TwitterStream(
 122         auth=auth,
 123         domain='userstream.twitter.com',
 124         **stream_args)
 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'):
 138             line1, line2, line3, line4 = draw(t=tweet)
 139             printNicely(line1)
 140             printNicely(line2)
 141             printNicely(line3)
 142             printNicely(line4)