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