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