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     parser = argparse.ArgumentParser(description=__doc__ or "")
  72     parser.add_argument(
  73         '-to',
  74         '--timeout',
  75         help='Timeout for the stream (seconds).')
  76     parser.add_argument(
  77         '-ht',
  78         '--heartbeat-timeout',
  79         help='Set heartbeat timeout.',
  80         default=90)
  81     parser.add_argument(
  82         '-nb',
  83         '--no-block',
  84         action='store_true',
  85         help='Set stream to non-blocking.')
  86     parser.add_argument(
  87         '-tt',
  88         '--track-keywords',
  89         help='Search the stream for specific text.')
  90     return parser.parse_args()
  91
  92
  93 def authen():
  94     """
  95     authenticate with Twitter OAuth
  96     """
  97     # When using rainbow stream you must authorize.
  98     twitter_credential = os.environ.get(
  99         'HOME',
 100         os.environ.get(
 101             'USERPROFILE',
 102             '')) + os.sep + '.rainbow_oauth'
 103     if not os.path.exists(twitter_credential):
 104         oauth_dance("Rainbow Stream",
 105                     CONSUMER_KEY,
 106                     CONSUMER_SECRET,
 107                     twitter_credential)
 108     oauth_token, oauth_token_secret = read_token_file(twitter_credential)
 109     return OAuth(
 110         oauth_token,
 111         oauth_token_secret,
 112         CONSUMER_KEY,
 113         CONSUMER_SECRET)
 114
 115
 116 def get_decorated_name():
 117     """
 118     Beginning of every line
 119     """
 120     t = Twitter(auth=authen())
 121     name = '@' + t.statuses.user_timeline()[-1]['user']['screen_name']
 122     g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
 123
 124
 125 def tweet():
 126     """
 127     Authen and tweet
 128     """
 129     t = Twitter(auth=authen())
 130     t.statuses.update(status=g['stuff'])
 131
 132
 133 def search():
 134     """
 135     Authen and search
 136     """
 137     t = Twitter(auth=authen())
 138     rel = t.search.tweets(q='#' + g['stuff'])['statuses']
 139     h, w = os.popen('stty size', 'r').read().split()
 140
 141     printNicely(grey('*'*int(w)+'\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('*'*int(w)+'\n'))
 146
 147
 148 def friend():
 149     """
 150     List of friend (following)
 151     """
 152     t = Twitter(auth=authen())
 153     g['friends'] = t.friends.ids()['ids']
 154     for i in g['friends']:
 155         screen_name = t.users.lookup(user_id=i)[0]['screen_name']
 156         user = cycle_color('@'+screen_name)
 157         print(user, end=' ')
 158     print('\n');
 159
 160
 161 def follower():
 162     """
 163     List of follower
 164     """
 165     t = Twitter(auth=authen())
 166     g['followers'] = t.followers.ids()['ids']
 167     for i in g['followers']:
 168         screen_name = t.users.lookup(user_id=i)[0]['screen_name']
 169         user = cycle_color('@'+screen_name)
 170         print(user, end=' ')
 171     print('\n');
 172
 173
 174 def help():
 175     """
 176     Print help
 177     """
 178     usage = '''
 179     Hi boss! I'm ready to serve you right now!
 180     ----------------------------------------------------
 181     "tweet" at the beginning will tweet immediately
 182     "s" and follow by any word will search and return 5 newest tweet
 183     "fr" will list out your following people
 184     "fl" will list out your followers
 185     "h" will print this help once again
 186     "c" will clear the terminal
 187     "q" will exit
 188     ----------------------------------------------------
 189     Have fun and hang tight!
 190     '''
 191     printNicely(usage)
 192     sys.stdout.write(g['decorated_name'])
 193
 194
 195 def clear():
 196     """
 197     Exit all
 198     """
 199     os.system('clear')
 200
 201
 202 def quit():
 203     """
 204     Exit all
 205     """
 206     os.kill(g['stream_pid'], signal.SIGKILL)
 207     sys.exit()
 208
 209
 210 def process(cmd):
 211     """
 212     Process switch by start of line
 213     """
 214     return {
 215         'tweet' : tweet,
 216         's'     : search,
 217         'fr'    : friend,
 218         'fl'    : follower,
 219         'h'     : help,
 220         'c'     : clear,
 221         'q'     : quit,
 222     }.get(cmd,lambda: sys.stdout.write(g['decorated_name']))
 223
 224
 225 def listen(stdin):
 226     """
 227     Listen to user's input
 228     """
 229     for line in iter(stdin.readline, ''):
 230         try:
 231             cmd = line.split()[0]
 232         except:
 233             cmd = ''
 234         # Save cmd to global variable and call process
 235         g['stuff'] = ' '.join(line.split()[1:])
 236         process(cmd)()
 237     stdin.close()
 238
 239
 240 def stream():
 241     """
 242     Track the stream
 243     """
 244     args = parse_arguments()
 245
 246     # The Logo
 247     ascii_art()
 248     # These arguments are optional:
 249     stream_args = dict(
 250         timeout=args.timeout,
 251         block=not args.no_block,
 252         heartbeat_timeout=args.heartbeat_timeout)
 253
 254     # Track keyword
 255     query_args = dict()
 256     if args.track_keywords:
 257         query_args['track'] = args.track_keywords
 258
 259     # Get stream
 260     stream = TwitterStream(
 261         auth = authen(),
 262         domain = 'userstream.twitter.com',
 263         **stream_args)
 264     tweet_iter = stream.user(**query_args)
 265
 266     # Iterate over the sample stream.
 267     for tweet in tweet_iter:
 268         if tweet is None:
 269             printNicely("-- None --")
 270         elif tweet is Timeout:
 271             printNicely("-- Timeout --")
 272         elif tweet is HeartbeatTimeout:
 273             printNicely("-- Heartbeat Timeout --")
 274         elif tweet is Hangup:
 275             printNicely("-- Hangup --")
 276         elif tweet.get('text'):
 277             draw(t=tweet)
 278
 279
 280 def fly():
 281     """
 282     Main function
 283     """
 284     get_decorated_name()
 285     p = Process(target = stream)
 286     p.start()
 287     g['stream_pid'] = p.pid
 288     listen(sys.stdin)
 289