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, imgflg = 0, 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 imgflg and 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     parser.add_argument(
 184         '-img',
 185         '--image',
 186         help='Display all photo on terminal.')
 187     return parser.parse_args()
 188
 189
 190 def authen():
 191     """
 192     Authenticate with Twitter OAuth
 193     """
 194     # When using rainbow stream you must authorize.
 195     twitter_credential = os.environ.get(
 196         'HOME',
 197         os.environ.get(
 198             'USERPROFILE',
 199             '')) + os.sep + '.rainbow_oauth'
 200     if not os.path.exists(twitter_credential):
 201         oauth_dance("Rainbow Stream",
 202                     CONSUMER_KEY,
 203                     CONSUMER_SECRET,
 204                     twitter_credential)
 205     oauth_token, oauth_token_secret = read_token_file(twitter_credential)
 206     return OAuth(
 207         oauth_token,
 208         oauth_token_secret,
 209         CONSUMER_KEY,
 210         CONSUMER_SECRET)
 211
 212
 213 def get_decorated_name():
 214     """
 215     Beginning of every line
 216     """
 217     t = Twitter(auth=authen())
 218     name = '@' + t.account.verify_credentials()['screen_name']
 219     g['original_name'] = name[1:]
 220     g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
 221
 222
 223 def switch():
 224     """
 225     Switch stream
 226     """
 227     try:
 228         target = g['stuff'].split()[0]
 229
 230         # Filter and ignore
 231         args = parse_arguments()
 232         try:
 233             if g['stuff'].split()[-1] == '-f':
 234                 only = raw_input('Only nicks: ')
 235                 ignore = raw_input('Ignore nicks: ')
 236                 args.filter = filter(None, only.split(','))
 237                 args.ignore = filter(None, ignore.split(','))
 238             elif g['stuff'].split()[-1] == '-d':
 239                 args.filter = ONLY_LIST
 240                 args.ignore = IGNORE_LIST
 241         except:
 242             printNicely(red('Sorry, wrong format.'))
 243             return
 244
 245         # Public stream
 246         if target == 'public':
 247             keyword = g['stuff'].split()[1]
 248             if keyword[0] == '#':
 249                 keyword = keyword[1:]
 250             # Kill old process
 251             os.kill(g['stream_pid'], signal.SIGKILL)
 252             args.track_keywords = keyword
 253             # Start new process
 254             p = Process(
 255                 target=stream,
 256                 args=(
 257                     PUBLIC_DOMAIN,
 258                     args))
 259             p.start()
 260             g['stream_pid'] = p.pid
 261
 262         # Personal stream
 263         elif target == 'mine':
 264             # Kill old process
 265             os.kill(g['stream_pid'], signal.SIGKILL)
 266             # Start new process
 267             p = Process(
 268                 target=stream,
 269                 args=(
 270                     USER_DOMAIN,
 271                     args,
 272                     g['original_name']))
 273             p.start()
 274             g['stream_pid'] = p.pid
 275         printNicely('')
 276         printNicely(green('Stream switched.'))
 277         if args.filter:
 278             printNicely(cyan('Only: ' + str(args.filter)))
 279         if args.ignore:
 280             printNicely(red('Ignore: ' + str(args.ignore)))
 281         printNicely('')
 282     except:
 283         printNicely(red('Sorry I can\'t understand.'))
 284
 285
 286 def home():
 287     """
 288     Home
 289     """
 290     t = Twitter(auth=authen())
 291     num = HOME_TWEET_NUM
 292     if g['stuff'].isdigit():
 293         num = g['stuff']
 294     for tweet in reversed(t.statuses.home_timeline(count=num)):
 295         draw(t=tweet, imgflg=g['image'])
 296     printNicely('')
 297
 298
 299 def view():
 300     """
 301     Friend view
 302     """
 303     t = Twitter(auth=authen())
 304     user = g['stuff'].split()[0]
 305     if user[0] == '@':
 306         try:
 307             num = int(g['stuff'].split()[1])
 308         except:
 309             num = HOME_TWEET_NUM
 310         for tweet in reversed(t.statuses.user_timeline(count=num, screen_name=user[1:])):
 311             draw(t=tweet, imgflg=g['image'])
 312         printNicely('')
 313     else:
 314         printNicely(red('A name should begin with a \'@\''))
 315
 316
 317 def tweet():
 318     """
 319     Tweet
 320     """
 321     t = Twitter(auth=authen())
 322     t.statuses.update(status=g['stuff'])
 323
 324
 325 def retweet():
 326     """
 327     ReTweet
 328     """
 329     t = Twitter(auth=authen())
 330     try:
 331         id = int(g['stuff'].split()[0])
 332         tid = db.rainbow_query(id)[0].tweet_id
 333         t.statuses.retweet(id=tid, include_entities=False, trim_user=True)
 334     except:
 335         printNicely(red('Sorry I can\'t retweet for you.'))
 336
 337
 338 def favorite():
 339     """
 340     Favorite
 341     """
 342     t = Twitter(auth=authen())
 343     try:
 344         id = int(g['stuff'].split()[0])
 345         tid = db.rainbow_query(id)[0].tweet_id
 346         t.favorites.create(_id=tid, include_entities=False)
 347         printNicely(green('Favorited.'))
 348         draw(t.statuses.show(id=tid), imgflg=g['image'])
 349     except:
 350         printNicely(red('Omg some syntax is wrong.'))
 351
 352
 353 def reply():
 354     """
 355     Reply
 356     """
 357     t = Twitter(auth=authen())
 358     try:
 359         id = int(g['stuff'].split()[0])
 360         tid = db.rainbow_query(id)[0].tweet_id
 361         user = t.statuses.show(id=tid)['user']['screen_name']
 362         status = ' '.join(g['stuff'].split()[1:])
 363         status = '@' + user + ' ' + status.decode('utf-8')
 364         t.statuses.update(status=status, in_reply_to_status_id=tid)
 365     except:
 366         printNicely(red('Sorry I can\'t understand.'))
 367
 368
 369 def delete():
 370     """
 371     Delete
 372     """
 373     t = Twitter(auth=authen())
 374     try:
 375         id = int(g['stuff'].split()[0])
 376         tid = db.rainbow_query(id)[0].tweet_id
 377         t.statuses.destroy(id=tid)
 378         printNicely(green('Okay it\'s gone.'))
 379     except:
 380         printNicely(red('Sorry I can\'t delete this tweet for you.'))
 381
 382
 383 def unfavorite():
 384     """
 385     Unfavorite
 386     """
 387     t = Twitter(auth=authen())
 388     try:
 389         id = int(g['stuff'].split()[0])
 390         tid = db.rainbow_query(id)[0].tweet_id
 391         t.favorites.destroy(_id=tid)
 392         printNicely(green('Okay it\'s unfavorited.'))
 393         draw(t.statuses.show(id=tid), imgflg=g['image'])
 394     except:
 395         printNicely(red('Sorry I can\'t unfavorite this tweet for you.'))
 396
 397
 398 def search():
 399     """
 400     Search
 401     """
 402     t = Twitter(auth=authen())
 403     try:
 404         if g['stuff'][0] == '#':
 405             rel = t.search.tweets(q=g['stuff'])['statuses']
 406             if len(rel):
 407                 printNicely('Newest tweets:')
 408                 for i in reversed(xrange(SEARCH_MAX_RECORD)):
 409                     draw(t=rel[i],
 410                         imgflg=g['image'],
 411                         keyword=g['stuff'].strip()[1:])
 412                 printNicely('')
 413             else:
 414                 printNicely(magenta('I\'m afraid there is no result'))
 415         else:
 416             printNicely(red('A keyword should be a hashtag (like \'#AKB48\')'))
 417     except:
 418         printNicely(red('Sorry I can\'t understand.'))
 419
 420
 421 def friend():
 422     """
 423     List of friend (following)
 424     """
 425     t = Twitter(auth=authen())
 426     g['friends'] = t.friends.ids()['ids']
 427     for i in g['friends']:
 428         name = t.users.lookup(user_id=i)[0]['name']
 429         screen_name = '@' + t.users.lookup(user_id=i)[0]['screen_name']
 430         user = cycle_color(name) + grey(' ' + screen_name + ' ')
 431         print user
 432
 433
 434 def follower():
 435     """
 436     List of follower
 437     """
 438     t = Twitter(auth=authen())
 439     g['followers'] = t.followers.ids()['ids']
 440     for i in g['followers']:
 441         name = t.users.lookup(user_id=i)[0]['name']
 442         screen_name = '@' + t.users.lookup(user_id=i)[0]['screen_name']
 443         user = cycle_color(name) + grey(' ' + screen_name + ' ')
 444         print user
 445
 446
 447 def help():
 448     """
 449     Help
 450     """
 451     s = ' ' * 2
 452     h, w = os.popen('stty size', 'r').read().split()
 453
 454     usage = '\n'
 455     usage += s + 'Hi boss! I\'m ready to serve you right now!\n'
 456     usage += s + '-' * (int(w) - 4) + '\n'
 457
 458     usage += s + 'You are ' + yellow('already') + ' on your personal stream.\n'
 459     usage += s * 2 + green('switch public #AKB') + \
 460         ' will switch to public stream and follow "' + \
 461         yellow('AKB') + '" keyword.\n'
 462     usage += s * 2 + green('switch mine') + \
 463         ' will switch to your personal stream.\n'
 464     usage += s * 2 + green('switch mine -f ') + \
 465         ' will prompt to enter the filter.\n'
 466     usage += s * 3 + yellow('Only nicks') + \
 467         ' filter will decide nicks will be INCLUDE ONLY.\n'
 468     usage += s * 3 + yellow('Ignore nicks') + \
 469         ' filter will decide nicks will be EXCLUDE.\n'
 470     usage += s * 2 + green('switch mine -d') + \
 471         ' will use the config\'s ONLY_LIST and IGNORE_LIST.\n'
 472     usage += s * 3 + '(see ' + grey('rainbowstream/config.py') + ').\n'
 473
 474     usage += s + 'For more action: \n'
 475     usage += s * 2 + green('home') + ' will show your timeline. ' + \
 476         green('home 7') + ' will show 7 tweet.\n'
 477     usage += s * 2 + green('view @mdo') + \
 478         ' will show ' + yellow('@mdo') + '\'s home.\n'
 479     usage += s * 2 + green('t oops ') + \
 480         'will tweet "' + yellow('oops') + '" immediately.\n'
 481     usage += s * 2 + \
 482         green('rt 12 ') + ' will retweet to tweet with ' + \
 483         yellow('[id=12]') + '.\n'
 484     usage += s * 2 + \
 485         green('fav 12 ') + ' will favorite the tweet with ' + \
 486         yellow('[id=12]') + '.\n'
 487     usage += s * 2 + green('rep 12 oops') + ' will reply "' + \
 488         yellow('oops') + '" to tweet with ' + yellow('[id=12]') + '.\n'
 489     usage += s * 2 + \
 490         green('del 12 ') + ' will delete tweet with ' + \
 491         yellow('[id=12]') + '.\n'
 492     usage += s * 2 + \
 493         green('ufav 12 ') + ' will unfavorite tweet with ' + \
 494         yellow('[id=12]') + '.\n'
 495     usage += s * 2 + green('s #AKB48') + ' will search for "' + \
 496         yellow('AKB48') + '" and return 5 newest tweet.\n'
 497     usage += s * 2 + green('fr') + ' will list out your following people.\n'
 498     usage += s * 2 + green('fl') + ' will list out your follower.\n'
 499     usage += s * 2 + green('h') + ' will show this help again.\n'
 500     usage += s * 2 + green('c') + ' will clear the screen.\n'
 501     usage += s * 2 + green('q') + ' will quit.\n'
 502
 503     usage += s + '-' * (int(w) - 4) + '\n'
 504     usage += s + 'Have fun and hang tight!\n'
 505     printNicely(usage)
 506
 507
 508 def clear():
 509     """
 510     Clear screen
 511     """
 512     os.system('clear')
 513
 514
 515 def quit():
 516     """
 517     Exit all
 518     """
 519     os.system('rm -rf rainbow.db')
 520     os.kill(g['stream_pid'], signal.SIGKILL)
 521     sys.exit()
 522
 523
 524 def reset():
 525     """
 526     Reset prefix of line
 527     """
 528     if g['reset']:
 529         printNicely(magenta('Need tips ? Type "h" and hit Enter key!'))
 530     g['reset'] = False
 531
 532
 533 def process(cmd):
 534     """
 535     Process switch
 536     """
 537     return dict(zip(
 538         cmdset,
 539         [
 540             switch,
 541             home,
 542             view,
 543             tweet,
 544             retweet,
 545             favorite,
 546             reply,
 547             delete,
 548             unfavorite,
 549             search,
 550             friend,
 551             follower,
 552             help,
 553             clear,
 554             quit
 555         ]
 556     )).get(cmd, reset)
 557
 558
 559 def listen():
 560     """
 561     Listen to user's input
 562     """
 563     d = dict(zip(
 564         cmdset,
 565         [
 566             ['public #', 'mine'],  # switch
 567             [],  # home
 568             ['@'],  # view
 569             [],  # tweet
 570             [],  # retweet
 571             [],  # reply
 572             [],  # delete
 573             ['#'],  # search
 574             [],  # friend
 575             [],  # follower
 576             [],  # help
 577             [],  # clear
 578             [],  # quit
 579         ]
 580     ))
 581     init_interactive_shell(d)
 582     reset()
 583     while True:
 584         if g['prefix']:
 585             line = raw_input(g['decorated_name'])
 586         else:
 587             line = raw_input()
 588         try:
 589             cmd = line.split()[0]
 590         except:
 591             cmd = ''
 592         # Save cmd to global variable and call process
 593         g['stuff'] = ' '.join(line.split()[1:])
 594         process(cmd)()
 595         if cmd in ['switch', 't', 'rt', 'rep']:
 596             g['prefix'] = False
 597         else:
 598             g['prefix'] = True
 599
 600
 601 def stream(domain, args, name='Rainbow Stream'):
 602     """
 603     Track the stream
 604     """
 605
 606     # The Logo
 607     art_dict = {
 608         USER_DOMAIN: name,
 609         PUBLIC_DOMAIN: args.track_keywords,
 610         SITE_DOMAIN: 'Site Stream',
 611     }
 612     ascii_art(art_dict[domain])
 613
 614     # These arguments are optional:
 615     stream_args = dict(
 616         timeout=args.timeout,
 617         block=not args.no_block,
 618         heartbeat_timeout=args.heartbeat_timeout)
 619
 620     # Track keyword
 621     query_args = dict()
 622     if args.track_keywords:
 623         query_args['track'] = args.track_keywords
 624
 625     # Get stream
 626     stream = TwitterStream(
 627         auth=authen(),
 628         domain=domain,
 629         **stream_args)
 630
 631     if domain == USER_DOMAIN:
 632         tweet_iter = stream.user(**query_args)
 633     elif domain == SITE_DOMAIN:
 634         tweet_iter = stream.site(**query_args)
 635     else:
 636         if args.track_keywords:
 637             tweet_iter = stream.statuses.filter(**query_args)
 638         else:
 639             tweet_iter = stream.statuses.sample()
 640
 641     # Iterate over the stream.
 642     for tweet in tweet_iter:
 643         if tweet is None:
 644             printNicely("-- None --")
 645         elif tweet is Timeout:
 646             printNicely("-- Timeout --")
 647         elif tweet is HeartbeatTimeout:
 648             printNicely("-- Heartbeat Timeout --")
 649         elif tweet is Hangup:
 650             printNicely("-- Hangup --")
 651         elif tweet.get('text'):
 652             draw(
 653                 t=tweet,
 654                 imgflg=args.image,
 655                 keyword=args.track_keywords,
 656                 fil=args.filter,
 657                 ig=args.ignore,
 658                 )
 659
 660
 661 def fly():
 662     """
 663     Main function
 664     """
 665     # Spawn stream process
 666     args = parse_arguments()
 667     get_decorated_name()
 668     p = Process(target=stream, args=(USER_DOMAIN, args, g['original_name']))
 669     p.start()
 670
 671     # Start listen process
 672     time.sleep(0.5)
 673     g['reset'] = True
 674     g['prefix'] = True
 675     g['stream_pid'] = p.pid
 676     g['image'] = args.image
 677     listen()