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