rainbowstream/rainbow.py

   1 """
   2 Colorful user's timeline stream
   3 """
   4 from multiprocessing import Process
   5 from dateutil import parser
   6
   7 import os
   8 import os.path
   9 import sys
  10 import signal
  11 import argparse
  12 import time
  13 import datetime
  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
  21 from .colors import *
  22 from .config import *
  23 from .interactive import *
  24 from .db import *
  25
  26 g = {}
  27 db = RainbowDB()
  28 cmdset = [
  29     'switch',
  30     'home',
  31     'view',
  32     't',
  33     'rt',
  34     'rep',
  35     'del',
  36     's',
  37     'fr',
  38     'fl',
  39     'h',
  40     'c',
  41     'q'
  42 ]
  43
  44
  45 def draw(t, keyword=None, fil=[], ig=[]):
  46     """
  47     Draw the rainbow
  48     """
  49
  50     # Retrieve tweet
  51     tid = t['id']
  52     text = t['text']
  53     screen_name = t['user']['screen_name']
  54     name = t['user']['name']
  55     created_at = t['created_at']
  56     date = parser.parse(created_at)
  57     date = date - datetime.timedelta(seconds=time.timezone)
  58     clock = date.strftime('%Y/%m/%d %H:%M:%S')
  59
  60     # Filter and ignore
  61     screen_name = '@' + screen_name
  62     if fil and screen_name not in fil:
  63         return
  64     if ig and screen_name in ig:
  65         return
  66
  67     res = db.tweet_query(tid)
  68     if not res:
  69         db.store(tid)
  70         res = db.tweet_query(tid)
  71     rid = res[0].rainbow_id
  72
  73     # Format info
  74     user = cycle_color(name) + grey(' ' + screen_name + ' ')
  75     meta = grey('[' + clock + '] [id=' + str(rid) + ']')
  76     tweet = text.split()
  77     # Highlight RT
  78     tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
  79     # Highlight screen_name
  80     tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
  81     # Highlight link
  82     tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
  83     # Highlight search keyword
  84     if keyword:
  85         tweet = map(
  86             lambda x: on_yellow(x) if
  87             ''.join(c for c in x if c.isalnum()).lower() == keyword.lower()
  88             else x,
  89             tweet
  90         )
  91     tweet = ' '.join(tweet)
  92
  93     # Draw rainbow
  94     line1 = u"{u:>{uw}}:".format(
  95         u=user,
  96         uw=len(user) + 2,
  97     )
  98     line2 = u"{c:>{cw}}".format(
  99         c=meta,
 100         cw=len(meta) + 2,
 101     )
 102     line3 = '  ' + tweet
 103
 104     printNicely('')
 105     printNicely(line1)
 106     printNicely(line2)
 107     printNicely(line3)
 108
 109
 110 def parse_arguments():
 111     """
 112     Parse the arguments
 113     """
 114     parser = argparse.ArgumentParser(description=__doc__ or "")
 115     parser.add_argument(
 116         '-to',
 117         '--timeout',
 118         help='Timeout for the stream (seconds).')
 119     parser.add_argument(
 120         '-ht',
 121         '--heartbeat-timeout',
 122         help='Set heartbeat timeout.',
 123         default=90)
 124     parser.add_argument(
 125         '-nb',
 126         '--no-block',
 127         action='store_true',
 128         help='Set stream to non-blocking.')
 129     parser.add_argument(
 130         '-tt',
 131         '--track-keywords',
 132         help='Search the stream for specific text.')
 133     parser.add_argument(
 134         '-fil',
 135         '--filter',
 136         help='Filter specific screen_name.')
 137     parser.add_argument(
 138         '-ig',
 139         '--ignore',
 140         help='Ignore specific screen_name.')
 141     return parser.parse_args()
 142
 143
 144 def authen():
 145     """
 146     Authenticate with Twitter OAuth
 147     """
 148     # When using rainbow stream you must authorize.
 149     twitter_credential = os.environ.get(
 150         'HOME',
 151         os.environ.get(
 152             'USERPROFILE',
 153             '')) + os.sep + '.rainbow_oauth'
 154     if not os.path.exists(twitter_credential):
 155         oauth_dance("Rainbow Stream",
 156                     CONSUMER_KEY,
 157                     CONSUMER_SECRET,
 158                     twitter_credential)
 159     oauth_token, oauth_token_secret = read_token_file(twitter_credential)
 160     return OAuth(
 161         oauth_token,
 162         oauth_token_secret,
 163         CONSUMER_KEY,
 164         CONSUMER_SECRET)
 165
 166
 167 def get_decorated_name():
 168     """
 169     Beginning of every line
 170     """
 171     t = Twitter(auth=authen())
 172     name = '@' + t.account.verify_credentials()['screen_name']
 173     g['original_name'] = name[1:]
 174     g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
 175
 176
 177 def switch():
 178     """
 179     Switch stream
 180     """
 181     try:
 182         target = g['stuff'].split()[0]
 183
 184         # Filter and ignore
 185         args = parse_arguments()
 186         try :
 187             if g['stuff'].split()[-1] == '-f':
 188                 only = raw_input('Only nicks: ')
 189                 ignore = raw_input('Ignore nicks: ')
 190                 args.filter = filter(None,only.split(','))
 191                 args.ignore = filter(None,ignore.split(','))
 192             elif g['stuff'].split()[-1] == '-d':
 193                 args.filter = ONLY_LIST
 194                 args.ignore = IGNORE_LIST
 195         except:
 196             printNicely(red('Sorry, wrong format.'))
 197             return
 198
 199         # Public stream
 200         if target == 'public':
 201             keyword = g['stuff'].split()[1]
 202             if keyword[0] == '#':
 203                 keyword = keyword[1:]
 204             # Kill old process
 205             os.kill(g['stream_pid'], signal.SIGKILL)
 206             args.track_keywords = keyword
 207             # Start new process
 208             p = Process(
 209                 target=stream,
 210                 args=(
 211                     PUBLIC_DOMAIN,
 212                     args))
 213             p.start()
 214             g['stream_pid'] = p.pid
 215
 216         # Personal stream
 217         elif target == 'mine':
 218             # Kill old process
 219             os.kill(g['stream_pid'], signal.SIGKILL)
 220             # Start new process
 221             p = Process(
 222                 target=stream,
 223                 args=(
 224                     USER_DOMAIN,
 225                     args,
 226                     g['original_name']))
 227             p.start()
 228             g['stream_pid'] = p.pid
 229         printNicely('')
 230         printNicely(green('Stream switched.'))
 231         if args.filter:
 232             printNicely(cyan('Only: ' + str(args.filter)))
 233         if args.ignore:
 234             printNicely(red('Ignore: ' + str(args.ignore)))
 235         printNicely('')
 236     except:
 237         printNicely(red('Sorry I can\'t understand.'))
 238
 239
 240 def home():
 241     """
 242     Home
 243     """
 244     t = Twitter(auth=authen())
 245     num = HOME_TWEET_NUM
 246     if g['stuff'].isdigit():
 247         num = g['stuff']
 248     for tweet in reversed(t.statuses.home_timeline(count=num)):
 249         draw(t=tweet)
 250     printNicely('')
 251
 252
 253 def view():
 254     """
 255     Friend view
 256     """
 257     t = Twitter(auth=authen())
 258     user = g['stuff'].split()[0]
 259     if user[0] == '@':
 260         try:
 261             num = int(g['stuff'].split()[1])
 262         except:
 263             num = HOME_TWEET_NUM
 264         for tweet in reversed(t.statuses.user_timeline(count=num, screen_name=user[1:])):
 265             draw(t=tweet)
 266         printNicely('')
 267     else:
 268         printNicely(red('A name should begin with a \'@\''))
 269
 270
 271 def tweet():
 272     """
 273     Tweet
 274     """
 275     t = Twitter(auth=authen())
 276     t.statuses.update(status=g['stuff'])
 277
 278
 279 def retweet():
 280     """
 281     ReTweet
 282     """
 283     t = Twitter(auth=authen())
 284     try:
 285         id = int(g['stuff'].split()[0])
 286         tid = db.rainbow_query(id)[0].tweet_id
 287         t.statuses.retweet(id=tid, include_entities=False, trim_user=True)
 288     except:
 289         printNicely(red('Sorry I can\'t retweet for you.'))
 290
 291
 292 def reply():
 293     """
 294     Reply
 295     """
 296     t = Twitter(auth=authen())
 297     try:
 298         id = int(g['stuff'].split()[0])
 299         tid = db.rainbow_query(id)[0].tweet_id
 300         user = t.statuses.show(id=tid)['user']['screen_name']
 301         status = ' '.join(g['stuff'].split()[1:])
 302         status = '@' + user + ' ' + status.decode('utf-8')
 303         t.statuses.update(status=status, in_reply_to_status_id=tid)
 304     except:
 305         printNicely(red('Sorry I can\'t understand.'))
 306
 307
 308 def delete():
 309     """
 310     Delete
 311     """
 312     t = Twitter(auth=authen())
 313     try:
 314         id = int(g['stuff'].split()[0])
 315         tid = db.rainbow_query(id)[0].tweet_id
 316         t.statuses.destroy(id=tid)
 317         printNicely(green('Okay it\'s gone.'))
 318     except:
 319         printNicely(red('Sorry I can\'t delete this tweet for you.'))
 320
 321
 322 def search():
 323     """
 324     Search
 325     """
 326     t = Twitter(auth=authen())
 327     try:
 328         if g['stuff'][0] == '#':
 329             rel = t.search.tweets(q=g['stuff'])['statuses']
 330             if len(rel):
 331                 printNicely('Newest tweets:')
 332                 for i in reversed(xrange(SEARCH_MAX_RECORD)):
 333                     draw(t=rel[i], keyword=g['stuff'].strip()[1:])
 334                 printNicely('')
 335             else:
 336                 printNicely(magenta('I\'m afraid there is no result'))
 337         else:
 338             printNicely(red('A keyword should be a hashtag (like \'#AKB48\')'))
 339     except:
 340         printNicely(red('Sorry I can\'t understand.'))
 341
 342
 343 def friend():
 344     """
 345     List of friend (following)
 346     """
 347     t = Twitter(auth=authen())
 348     g['friends'] = t.friends.ids()['ids']
 349     for i in g['friends']:
 350         name = t.users.lookup(user_id=i)[0]['name']
 351         screen_name = '@' + t.users.lookup(user_id=i)[0]['screen_name']
 352         user = cycle_color(name) + grey(' ' + screen_name + ' ')
 353         print user
 354
 355
 356 def follower():
 357     """
 358     List of follower
 359     """
 360     t = Twitter(auth=authen())
 361     g['followers'] = t.followers.ids()['ids']
 362     for i in g['followers']:
 363         name = t.users.lookup(user_id=i)[0]['name']
 364         screen_name = '@' + t.users.lookup(user_id=i)[0]['screen_name']
 365         user = cycle_color(name) + grey(' ' + screen_name + ' ')
 366         print user
 367
 368
 369 def help():
 370     """
 371     Help
 372     """
 373     s = ' '*2
 374     h,w = os.popen('stty size', 'r').read().split()
 375
 376     usage = '\n'
 377     usage += s + 'Hi boss! I\'m ready to serve you right now!\n'
 378     usage += s + '-'*(int(w)-4) + '\n'
 379
 380     usage += s + 'You are ' + yellow('already') +' on your personal stream.\n'
 381     usage += s*2 + green('switch public #AKB') +  ' will switch to public stream and follow "' + yellow('AKB') + '" keyword.\n'
 382     usage += s*2 + green('switch mine') + ' will switch to your personal stream.\n'
 383     usage += s*2 + green('switch mine -f ') + ' will prompt to enter the filter.\n'
 384     usage += s*3 + yellow('Only nicks') + ' filter will decide nicks will be INCLUDE ONLY.\n'
 385     usage += s*3 + yellow('Ignore nicks') + ' filter will decide nicks will be EXCLUDE.\n'
 386     usage += s*2 + green('switch mine -d') + ' will use the config\'s ONLY_LIST and IGNORE_LIST.\n'
 387     usage += s*3 + '(see ' + grey('rainbowstream/config.py') + ').\n'
 388
 389     usage += s + 'For more action: \n'
 390     usage += s*2 + green('home') + ' will show your timeline. ' + green('home 7') + ' will show 7 tweet.\n'
 391     usage += s*2 + green('view @mdo') + ' will show ' + yellow('@mdo') + '\'s home.\n'
 392     usage += s*2 + green('t oops ') + 'will tweet "' + yellow('oops') + '" immediately.\n'
 393     usage += s*2 + green('rt 12 ') + ' will retweet to tweet with '+ yellow('[id=12]') + '.\n'
 394     usage += s*2 + green('rep 12 oops') + ' will reply "' + yellow('oops') + '" to tweet with ' + yellow('[id=12]') + '.\n'
 395     usage += s*2 + green('del 12 ') + ' will delete tweet with ' + yellow('[id=12]') + '.\n'
 396     usage += s*2 + green('s #AKB48') + ' will search for "'+yellow('AKB48') + '" and return 5 newest tweet.\n'
 397     usage += s*2 + green('fr') + ' will list out your following people.\n'
 398     usage += s*2 + green('fl') + ' will list out your follower.\n'
 399     usage += s*2 + green('h') + ' will show this help again.\n'
 400     usage += s*2 + green('c') + ' will clear the screen.\n'
 401     usage += s*2 + green('q') + ' will quit.\n'
 402
 403     usage += s + '-'*(int(w)-4) + '\n'
 404     usage += s + 'Have fun and hang tight!\n'
 405     printNicely(usage)
 406
 407
 408 def clear():
 409     """
 410     Clear screen
 411     """
 412     os.system('clear')
 413
 414
 415 def quit():
 416     """
 417     Exit all
 418     """
 419     os.system('rm -rf rainbow.db')
 420     os.kill(g['stream_pid'], signal.SIGKILL)
 421     sys.exit()
 422
 423
 424 def reset():
 425     """
 426     Reset prefix of line
 427     """
 428     if g['reset']:
 429         printNicely(magenta('Need tips ? Type "h" and hit Enter key!'))
 430     g['reset'] = False
 431
 432
 433 def process(cmd):
 434     """
 435     Process switch
 436     """
 437     return dict(zip(
 438         cmdset,
 439         [
 440             switch,
 441             home,
 442             view,
 443             tweet,
 444             retweet,
 445             reply,
 446             delete,
 447             search,
 448             friend,
 449             follower,
 450             help,
 451             clear,
 452             quit
 453         ]
 454     )).get(cmd, reset)
 455
 456
 457 def listen():
 458     """
 459     Listen to user's input
 460     """
 461     d = dict(zip(
 462         cmdset,
 463         [
 464             ['public #','mine'], # switch
 465             [], # home
 466             ['@'], # view
 467             [], # tweet
 468             [], # retweet
 469             [], # reply
 470             [], # delete
 471             ['#'], # search
 472             [], # friend
 473             [], # follower
 474             [], # help
 475             [], # clear
 476             [], # quit
 477         ]
 478         ))
 479     init_interactive_shell(d)
 480     reset()
 481     while True:
 482         if g['prefix']:
 483             line = raw_input(g['decorated_name'])
 484         else:
 485             line = raw_input()
 486         try:
 487             cmd = line.split()[0]
 488         except:
 489             cmd = ''
 490         # Save cmd to global variable and call process
 491         g['stuff'] = ' '.join(line.split()[1:])
 492         process(cmd)()
 493         if cmd in ['switch','t','rt','rep']:
 494             g['prefix'] = False
 495         else:
 496             g['prefix'] = True
 497
 498
 499 def stream(domain, args, name='Rainbow Stream'):
 500     """
 501     Track the stream
 502     """
 503
 504     # The Logo
 505     art_dict = {
 506         USER_DOMAIN: name,
 507         PUBLIC_DOMAIN: args.track_keywords,
 508         SITE_DOMAIN: 'Site Stream',
 509     }
 510     ascii_art(art_dict[domain])
 511
 512     # These arguments are optional:
 513     stream_args = dict(
 514         timeout=args.timeout,
 515         block=not args.no_block,
 516         heartbeat_timeout=args.heartbeat_timeout)
 517
 518     # Track keyword
 519     query_args = dict()
 520     if args.track_keywords:
 521         query_args['track'] = args.track_keywords
 522
 523     # Get stream
 524     stream = TwitterStream(
 525         auth=authen(),
 526         domain=domain,
 527         **stream_args)
 528
 529     if domain == USER_DOMAIN:
 530         tweet_iter = stream.user(**query_args)
 531     elif domain == SITE_DOMAIN:
 532         tweet_iter = stream.site(**query_args)
 533     else:
 534         if args.track_keywords:
 535             tweet_iter = stream.statuses.filter(**query_args)
 536         else:
 537             tweet_iter = stream.statuses.sample()
 538
 539     # Iterate over the stream.
 540     for tweet in tweet_iter:
 541         if tweet is None:
 542             printNicely("-- None --")
 543         elif tweet is Timeout:
 544             printNicely("-- Timeout --")
 545         elif tweet is HeartbeatTimeout:
 546             printNicely("-- Heartbeat Timeout --")
 547         elif tweet is Hangup:
 548             printNicely("-- Hangup --")
 549         elif tweet.get('text'):
 550             draw(t=tweet, keyword=args.track_keywords, fil=args.filter, ig=args.ignore)
 551
 552
 553 def fly():
 554     """
 555     Main function
 556     """
 557     # Spawn stream process
 558     args = parse_arguments()
 559     get_decorated_name()
 560     p = Process(target=stream, args=(USER_DOMAIN, args, g['original_name']))
 561     p.start()
 562
 563     # Start listen process
 564     time.sleep(0.5)
 565     g['reset'] = True
 566     g['prefix'] = True
 567     g['stream_pid'] = p.pid
 568     listen()
 569