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