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