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, os.path, argparse, sys, time
   9
  10 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
  11 from twitter.api import *
  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 auth_obj = {}
  21
  22 def draw(t):
  23     """
  24     Draw the rainbow
  25     """
  26     # Retrieve tweet
  27     text = t['text']
  28     screen_name = t['user']['screen_name']
  29     name = t['user']['name']
  30     created_at = t['created_at']
  31     date = parser.parse(created_at)
  32     time = date.strftime('%Y/%m/%d %H:%M:%S')
  33
  34     # Format info
  35     user = cycle_color(name) + grey(' ' + '@' + screen_name + ' ')
  36     clock = grey('[' + time + ']')
  37     tweet = text.split()
  38     tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
  39     tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
  40     tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
  41     tweet = ' '.join(tweet)
  42
  43     # Draw rainbow
  44     line1 = u"{u:>{uw}}:".format(
  45         u=user,
  46         uw=len(user) + 2,
  47     )
  48     line2 = u"{c:>{cw}}".format(
  49         c=clock,
  50         cw=len(clock) + 2,
  51     )
  52     line3 = '  ' + tweet
  53     line4 = '\n'
  54
  55     return line1, line2, line3, line4
  56
  57
  58 def parse_arguments():
  59     """
  60     Parse the arguments
  61     """
  62
  63     parser = argparse.ArgumentParser(description=__doc__ or "")
  64
  65     parser.add_argument(
  66         '-to',
  67         '--timeout',
  68         help='Timeout for the stream (seconds).')
  69     parser.add_argument(
  70         '-ht',
  71         '--heartbeat-timeout',
  72         help='Set heartbeat timeout.',
  73         default=90)
  74     parser.add_argument(
  75         '-nb',
  76         '--no-block',
  77         action='store_true',
  78         help='Set stream to non-blocking.')
  79     parser.add_argument(
  80         '-tt',
  81         '--track-keywords',
  82         help='Search the stream for specific text.')
  83     return parser.parse_args()
  84
  85
  86 def authen():
  87     """
  88     authenticate with Twitter OAuth
  89     """
  90     # When using rainbow stream you must authorize.
  91     twitter_credential = os.environ.get(
  92         'HOME',
  93         os.environ.get(
  94             'USERPROFILE',
  95             '')) + os.sep + '.rainbow_oauth'
  96     if not os.path.exists(twitter_credential):
  97         oauth_dance("Rainbow Stream",
  98                     CONSUMER_KEY,
  99                     CONSUMER_SECRET,
 100                     twitter_credential)
 101     oauth_token, oauth_token_secret = read_token_file(twitter_credential)
 102     return OAuth(
 103         oauth_token,
 104         oauth_token_secret,
 105         CONSUMER_KEY,
 106         CONSUMER_SECRET)
 107
 108
 109 def get_decorated_name():
 110     """
 111     Beginning of every line
 112     """
 113     t = Twitter(auth=authen())
 114     name = '@' + t.statuses.user_timeline()[-1]['user']['screen_name']
 115     auth_obj['decorated_name'] = grey('[') + grey(name) + grey(']: ')
 116
 117 def tweet(stuff):
 118     """
 119     Authen and tweet
 120     """
 121     t = Twitter(auth=authen())
 122     t.statuses.update(status=stuff)
 123
 124
 125 def listen(stdin):
 126     """
 127     Listen to user's input
 128     """
 129     for line in iter(stdin.readline, ''):
 130         # Public tweet
 131         if line.startswith('!'):
 132             tweet(line[1:])
 133         else:
 134             sys.stdout.write(auth_obj['decorated_name'])
 135     stdin.close()
 136
 137
 138 def stream():
 139     """
 140     Ouput the stream
 141     """
 142     args = parse_arguments()
 143
 144     # The Logo
 145     ascii_art()
 146     print("Tip: Press ENTER and put a '!' in the beginning to start composing a new tweet")
 147     print('\n')
 148     # These arguments are optional:
 149     stream_args = dict(
 150         timeout=args.timeout,
 151         block=not args.no_block,
 152         heartbeat_timeout=args.heartbeat_timeout)
 153
 154     # Track keyword
 155     query_args = dict()
 156     if args.track_keywords:
 157         query_args['track'] = args.track_keywords
 158
 159     # Get stream
 160     stream = TwitterStream(
 161         auth = authen(),
 162         domain = 'userstream.twitter.com',
 163         **stream_args)
 164     tweet_iter = stream.user(**query_args)
 165
 166     # Iterate over the sample stream.
 167     for tweet in tweet_iter:
 168         if tweet is None:
 169             printNicely("-- None --")
 170         elif tweet is Timeout:
 171             printNicely("-- Timeout --")
 172         elif tweet is HeartbeatTimeout:
 173             printNicely("-- Heartbeat Timeout --")
 174         elif tweet is Hangup:
 175             printNicely("-- Hangup --")
 176         elif tweet.get('text'):
 177             line1, line2, line3, line4 = draw(t=tweet)
 178             printNicely(line1)
 179             printNicely(line2)
 180             printNicely(line3)
 181             printNicely(line4)
 182
 183
 184 def fly():
 185     """
 186     Main function
 187     """
 188     get_decorated_name()
 189     p = Process(target = stream)
 190     p.start()
 191     listen(sys.stdin)
 192