rainbowstream/rainbow.py

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