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