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, signal
   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 g = {}
  21
  22 def draw(t,keyword=None):
  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     # Highlight RT
  39     tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
  40     # Highlight screen_name
  41     tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
  42     # Highlight link
  43     tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
  44     # Highlight search keyword
  45     if keyword:
  46         tweet = map(lambda x: on_yellow(x) if ''.join(c for c in x if c.isalnum()).lower() == keyword.lower() else x, tweet)
  47     tweet = ' '.join(tweet)
  48
  49     # Draw rainbow
  50     line1 = u"{u:>{uw}}:".format(
  51         u=user,
  52         uw=len(user) + 2,
  53     )
  54     line2 = u"{c:>{cw}}".format(
  55         c=clock,
  56         cw=len(clock) + 2,
  57     )
  58     line3 = '  ' + tweet
  59     line4 = '\n'
  60
  61     printNicely(line1)
  62     printNicely(line2)
  63     printNicely(line3)
  64     printNicely(line4)
  65
  66
  67 def parse_arguments():
  68     """
  69     Parse the arguments
  70     """
  71
  72     parser = argparse.ArgumentParser(description=__doc__ or "")
  73
  74     parser.add_argument(
  75         '-to',
  76         '--timeout',
  77         help='Timeout for the stream (seconds).')
  78     parser.add_argument(
  79         '-ht',
  80         '--heartbeat-timeout',
  81         help='Set heartbeat timeout.',
  82         default=90)
  83     parser.add_argument(
  84         '-nb',
  85         '--no-block',
  86         action='store_true',
  87         help='Set stream to non-blocking.')
  88     parser.add_argument(
  89         '-tt',
  90         '--track-keywords',
  91         help='Search the stream for specific text.')
  92     return parser.parse_args()
  93
  94
  95 def authen():
  96     """
  97     authenticate with Twitter OAuth
  98     """
  99     # When using rainbow stream you must authorize.
 100     twitter_credential = os.environ.get(
 101         'HOME',
 102         os.environ.get(
 103             'USERPROFILE',
 104             '')) + os.sep + '.rainbow_oauth'
 105     if not os.path.exists(twitter_credential):
 106         oauth_dance("Rainbow Stream",
 107                     CONSUMER_KEY,
 108                     CONSUMER_SECRET,
 109                     twitter_credential)
 110     oauth_token, oauth_token_secret = read_token_file(twitter_credential)
 111     return OAuth(
 112         oauth_token,
 113         oauth_token_secret,
 114         CONSUMER_KEY,
 115         CONSUMER_SECRET)
 116
 117
 118 def get_decorated_name():
 119     """
 120     Beginning of every line
 121     """
 122     t = Twitter(auth=authen())
 123     name = '@' + t.statuses.user_timeline()[-1]['user']['screen_name']
 124     g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
 125
 126
 127 def tweet():
 128     """
 129     Authen and tweet
 130     """
 131     t = Twitter(auth=authen())
 132     t.statuses.update(status=g['stuff'])
 133
 134
 135 def search():
 136     """
 137     Authen and search
 138     """
 139     t = Twitter(auth=authen())
 140     rel = t.search.tweets(q='#' + g['stuff'])['statuses']
 141     printNicely(grey('**************************************************************************************\n'))
 142     print('Newest',SEARCH_MAX_RECORD, 'tweet: \n')
 143     for i in xrange(5):
 144         draw(t=rel[i],keyword=g['stuff'].strip())
 145     printNicely(grey('**************************************************************************************\n'))
 146
 147
 148 def help():
 149     """
 150     Print help
 151     """
 152     usage = '''
 153     Hi boss! I'm ready to serve you right now!
 154     ----------------------------------------------------
 155     "!" at the beginning will start to tweet immediately
 156     "/" at the beginning will search for 5 newest tweet
 157     "?" or "h" will print this help once again
 158     "c" will clear the terminal
 159     "q" will exit
 160     ----------------------------------------------------
 161     Hvae fun and hang tight!
 162     '''
 163     printNicely(usage)
 164     sys.stdout.write(g['decorated_name'])
 165
 166
 167 def quit():
 168     """
 169     Exit all
 170     """
 171     os.kill(g['stream_pid'], signal.SIGKILL)
 172     sys.exit()
 173
 174
 175 def clear():
 176     """
 177     Exit all
 178     """
 179     os.system('clear')
 180
 181
 182 def process(line):
 183     """
 184     Process switch by start of line
 185     """
 186     return {
 187         '!' : tweet,
 188         '/' : search,
 189         '?' : help,
 190         'h' : help,
 191         'c' : clear,
 192         'q' : quit,
 193     }.get(line[0],lambda: sys.stdout.write(g['decorated_name']))
 194
 195
 196 def listen(stdin):
 197     """
 198     Listen to user's input
 199     """
 200     for line in iter(stdin.readline, ''):
 201         # Save cmd to global variable and call process
 202         g['stuff'] = line[1:]
 203         process(line)()
 204     stdin.close()
 205
 206
 207 def stream():
 208     """
 209     Track the stream
 210     """
 211     args = parse_arguments()
 212
 213     # The Logo
 214     ascii_art()
 215     # These arguments are optional:
 216     stream_args = dict(
 217         timeout=args.timeout,
 218         block=not args.no_block,
 219         heartbeat_timeout=args.heartbeat_timeout)
 220
 221     # Track keyword
 222     query_args = dict()
 223     if args.track_keywords:
 224         query_args['track'] = args.track_keywords
 225
 226     # Get stream
 227     stream = TwitterStream(
 228         auth = authen(),
 229         domain = 'userstream.twitter.com',
 230         **stream_args)
 231     tweet_iter = stream.user(**query_args)
 232
 233     # Iterate over the sample stream.
 234     for tweet in tweet_iter:
 235         if tweet is None:
 236             printNicely("-- None --")
 237         elif tweet is Timeout:
 238             printNicely("-- Timeout --")
 239         elif tweet is HeartbeatTimeout:
 240             printNicely("-- Heartbeat Timeout --")
 241         elif tweet is Hangup:
 242             printNicely("-- Hangup --")
 243         elif tweet.get('text'):
 244             draw(t=tweet)
 245
 246
 247 def fly():
 248     """
 249     Main function
 250     """
 251     get_decorated_name()
 252     p = Process(target = stream)
 253     p.start()
 254     g['stream_pid'] = p.pid
 255     listen(sys.stdin)
 256