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
   9 import os.path
  10 import argparse
  11 import sys
  12 import time
  13 import signal
  14
  15 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
  16 from twitter.api import *
  17 from twitter.oauth import OAuth, read_token_file
  18 from twitter.oauth_dance import oauth_dance
  19 from twitter.util import printNicely
  20 from dateutil import parser
  21
  22 from .colors import *
  23 from .config import *
  24
  25 g = {}
  26
  27
  28 def draw(t, keyword=None):
  29     """
  30     Draw the rainbow
  31     """
  32     # Retrieve tweet
  33     tid = t['id']
  34     text = t['text']
  35     screen_name = t['user']['screen_name']
  36     name = t['user']['name']
  37     created_at = t['created_at']
  38     date = parser.parse(created_at)
  39     time = date.strftime('%Y/%m/%d %H:%M:%S')
  40
  41     # Format info
  42     user = cycle_color(name) + grey(' ' + '@' + screen_name + ' ')
  43     meta = grey('[' + time + '] [id=' + str(tid) + ']')
  44     tweet = text.split()
  45     # Highlight RT
  46     tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
  47     # Highlight screen_name
  48     tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
  49     # Highlight link
  50     tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
  51     # Highlight search keyword
  52     if keyword:
  53         tweet = map(
  54             lambda x: on_yellow(x) if
  55             ''.join(c for c in x if c.isalnum()).lower() == keyword.lower()
  56             else x,
  57             tweet
  58         )
  59     tweet = ' '.join(tweet)
  60
  61     # Draw rainbow
  62     line1 = u"{u:>{uw}}:".format(
  63         u=user,
  64         uw=len(user) + 2,
  65     )
  66     line2 = u"{c:>{cw}}".format(
  67         c=meta,
  68         cw=len(meta) + 2,
  69     )
  70     line3 = '  ' + tweet
  71
  72     printNicely(line1)
  73     printNicely(line2)
  74     printNicely(line3)
  75     printNicely('')
  76
  77
  78 def parse_arguments():
  79     """
  80     Parse the arguments
  81     """
  82     parser = argparse.ArgumentParser(description=__doc__ or "")
  83     parser.add_argument(
  84         '-to',
  85         '--timeout',
  86         help='Timeout for the stream (seconds).')
  87     parser.add_argument(
  88         '-ht',
  89         '--heartbeat-timeout',
  90         help='Set heartbeat timeout.',
  91         default=90)
  92     parser.add_argument(
  93         '-nb',
  94         '--no-block',
  95         action='store_true',
  96         help='Set stream to non-blocking.')
  97     parser.add_argument(
  98         '-tt',
  99         '--track-keywords',
 100         help='Search the stream for specific text.')
 101     return parser.parse_args()
 102
 103
 104 def authen():
 105     """
 106     Authenticate with Twitter OAuth
 107     """
 108     # When using rainbow stream you must authorize.
 109     twitter_credential = os.environ.get(
 110         'HOME',
 111         os.environ.get(
 112             'USERPROFILE',
 113             '')) + os.sep + '.rainbow_oauth'
 114     if not os.path.exists(twitter_credential):
 115         oauth_dance("Rainbow Stream",
 116                     CONSUMER_KEY,
 117                     CONSUMER_SECRET,
 118                     twitter_credential)
 119     oauth_token, oauth_token_secret = read_token_file(twitter_credential)
 120     return OAuth(
 121         oauth_token,
 122         oauth_token_secret,
 123         CONSUMER_KEY,
 124         CONSUMER_SECRET)
 125
 126
 127 def get_decorated_name():
 128     """
 129     Beginning of every line
 130     """
 131     t = Twitter(auth=authen())
 132     name = '@' + t.statuses.user_timeline()[-1]['user']['screen_name']
 133     g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
 134
 135
 136 def home():
 137     """
 138     Home
 139     """
 140     t = Twitter(auth=authen())
 141     count = HOME_TWEET_NUM
 142     if g['stuff'].isdigit():
 143         count = g['stuff']
 144     for tweet in reversed(t.statuses.home_timeline(count=count)):
 145         draw(t=tweet)
 146
 147
 148 def view():
 149     """
 150     Friend view
 151     """
 152     t = Twitter(auth=authen())
 153     user = g['stuff'].split()[0]
 154     try:
 155         count = int(g['stuff'].split()[1])
 156     except:
 157         count = HOME_TWEET_NUM
 158     for tweet in reversed(t.statuses.user_timeline(count=count, screen_name=user)):
 159         draw(t=tweet)
 160
 161
 162 def tweet():
 163     """
 164     Tweet
 165     """
 166     t = Twitter(auth=authen())
 167     t.statuses.update(status=g['stuff'])
 168
 169
 170 def reply():
 171     """
 172     Reply
 173     """
 174     t = Twitter(auth=authen())
 175     try:
 176         id = int(g['stuff'].split()[0])
 177         user = t.statuses.show(id=id)['user']['screen_name']
 178         status = ' '.join(g['stuff'].split()[1:])
 179         status = '@' + user + ' ' + status.decode('utf-8')
 180         t.statuses.update(status=status, in_reply_to_status_id=id)
 181     except:
 182         print(red('Sorry I can\'t understand.'))
 183
 184
 185 def delete():
 186     """
 187     Delete
 188     """
 189     t = Twitter(auth=authen())
 190     try:
 191         id = int(g['stuff'].split()[0])
 192         t.statuses.destroy(id=id)
 193         print(green('Okay it\'s gone.'))
 194     except:
 195         print(red('Sorry I can\'t delete this tweet for you.'))
 196
 197
 198 def search():
 199     """
 200     Search
 201     """
 202     t = Twitter(auth=authen())
 203     rel = t.search.tweets(q='#' + g['stuff'])['statuses']
 204     h, w = os.popen('stty size', 'r').read().split()
 205
 206     printNicely(grey('*' * int(w) + '\n'))
 207     print('Newest', SEARCH_MAX_RECORD, 'tweet: \n')
 208     for i in xrange(5):
 209         draw(t=rel[i], keyword=g['stuff'].strip())
 210     printNicely(grey('*' * int(w) + '\n'))
 211
 212
 213 def friend():
 214     """
 215     List of friend (following)
 216     """
 217     t = Twitter(auth=authen())
 218     g['friends'] = t.friends.ids()['ids']
 219     for i in g['friends']:
 220         screen_name = t.users.lookup(user_id=i)[0]['screen_name']
 221         user = cycle_color('@' + screen_name)
 222         print(user, end=' ')
 223     print('\n')
 224
 225
 226 def follower():
 227     """
 228     List of follower
 229     """
 230     t = Twitter(auth=authen())
 231     g['followers'] = t.followers.ids()['ids']
 232     for i in g['followers']:
 233         screen_name = t.users.lookup(user_id=i)[0]['screen_name']
 234         user = cycle_color('@' + screen_name)
 235         print(user, end=' ')
 236     print('\n')
 237
 238
 239 def help():
 240     """
 241     Help
 242     """
 243     usage = '''
 244     Hi boss! I'm ready to serve you right now!
 245     ----------------------------------------------------
 246     "home" will show your timeline. "home 7" will print 7 tweet.
 247     "view bob" will show your friend @bob's home.
 248     "t oops" will tweet "oops" immediately.
 249     "rep 12345 oops" will reply "oops" to tweet with id "12345".
 250     "del 12345" will delete tweet with id "12345".
 251     "s AKB48" will search for "AKB48" and return 5 newest tweet.
 252     "fr" will list out your following people.
 253     "fl" will list out your followers.
 254     "h" or "help" will print this help once again.
 255     "c" will clear the terminal.
 256     "q" will exit.
 257     ----------------------------------------------------
 258     Have fun and hang tight!
 259     '''
 260     printNicely(usage)
 261     sys.stdout.write(g['decorated_name'])
 262
 263
 264 def clear():
 265     """
 266     Clear screen
 267     """
 268     os.system('clear')
 269
 270
 271 def quit():
 272     """
 273     Exit all
 274     """
 275     os.kill(g['stream_pid'], signal.SIGKILL)
 276     sys.exit()
 277
 278
 279 def process(cmd):
 280     """
 281     Process switch
 282     """
 283     return {
 284         'home': home,
 285         'view': view,
 286         't': tweet,
 287         'rep': reply,
 288         'del': delete,
 289         's': search,
 290         'fr': friend,
 291         'fl': follower,
 292         'h': help,
 293         'help': help,
 294         'c': clear,
 295         'q': quit,
 296     }.get(cmd, lambda: sys.stdout.write(g['decorated_name']))
 297
 298
 299 def listen(stdin):
 300     """
 301     Listen to user's input
 302     """
 303     for line in iter(stdin.readline, ''):
 304         try:
 305             cmd = line.split()[0]
 306         except:
 307             cmd = ''
 308         # Save cmd to global variable and call process
 309         g['stuff'] = ' '.join(line.split()[1:])
 310         process(cmd)()
 311     stdin.close()
 312
 313
 314 def stream():
 315     """
 316     Track the stream
 317     """
 318     args = parse_arguments()
 319
 320     # The Logo
 321     ascii_art()
 322     # These arguments are optional:
 323     stream_args = dict(
 324         timeout=args.timeout,
 325         block=not args.no_block,
 326         heartbeat_timeout=args.heartbeat_timeout)
 327
 328     # Track keyword
 329     query_args = dict()
 330     if args.track_keywords:
 331         query_args['track'] = args.track_keywords
 332
 333     # Get stream
 334     stream = TwitterStream(
 335         auth=authen(),
 336         domain='userstream.twitter.com',
 337         **stream_args)
 338     tweet_iter = stream.user(**query_args)
 339
 340     # Iterate over the sample stream.
 341     for tweet in tweet_iter:
 342         if tweet is None:
 343             printNicely("-- None --")
 344         elif tweet is Timeout:
 345             printNicely("-- Timeout --")
 346         elif tweet is HeartbeatTimeout:
 347             printNicely("-- Heartbeat Timeout --")
 348         elif tweet is Hangup:
 349             printNicely("-- Hangup --")
 350         elif tweet.get('text'):
 351             draw(t=tweet)
 352
 353
 354 def fly():
 355     """
 356     Main function
 357     """
 358     get_decorated_name()
 359     p = Process(target=stream)
 360     p.start()
 361     g['stream_pid'] = p.pid
 362     listen(sys.stdin)