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