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(',')[0]
 192                 args.ignore = ignore.split(',')[0]
 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
 384         (see rainbowstream/config.py).
 385     For more action:
 386       "home" will show your timeline. "home 7" will show 7 tweet.
 387       "view @bob" will show your friend @bob's home.
 388       "t oops" will tweet "oops" immediately.
 389       "rt 12345" will retweet to tweet with id "12345".
 390       "rep 12345 oops" will reply "oops" to tweet with id "12345".
 391       "del 12345" will delete tweet with id "12345".
 392       "s #AKB48" will search for "AKB48" and return 5 newest tweet.
 393       "fr" will list out your following people.
 394       "fl" will list out your followers.
 395       "h" will show this help again.
 396       "c" will clear the terminal.
 397       "q" will exit.
 398     -------------------------------------------------------------
 399     Have fun and hang tight!
 400     '''
 401     printNicely(usage)
 402
 403
 404 def clear():
 405     """
 406     Clear screen
 407     """
 408     os.system('clear')
 409
 410
 411 def quit():
 412     """
 413     Exit all
 414     """
 415     os.system('rm -rf rainbow.db')
 416     os.kill(g['stream_pid'], signal.SIGKILL)
 417     sys.exit()
 418
 419
 420 def reset():
 421     """
 422     Reset prefix of line
 423     """
 424     if g['reset']:
 425         printNicely(green('Need tips ? Type "h" and hit Enter key!'))
 426     g['reset'] = False
 427
 428
 429 def process(cmd):
 430     """
 431     Process switch
 432     """
 433     return dict(zip(
 434         cmdset,
 435         [
 436             switch,
 437             home,
 438             view,
 439             tweet,
 440             retweet,
 441             reply,
 442             delete,
 443             search,
 444             friend,
 445             follower,
 446             help,
 447             clear,
 448             quit
 449         ]
 450     )).get(cmd, reset)
 451
 452
 453 def listen():
 454     """
 455     Listen to user's input
 456     """
 457     d = dict(zip(
 458         cmdset,
 459         [
 460             ['public #','mine'], # switch
 461             [], # home
 462             ['@'], # view
 463             [], # tweet
 464             [], # retweet
 465             [], # reply
 466             [], # delete
 467             ['#'], # search
 468             [], # friend
 469             [], # follower
 470             [], # help
 471             [], # clear
 472             [], # quit
 473         ]
 474         ))
 475     init_interactive_shell(d)
 476     reset()
 477     while True:
 478         if g['prefix']:
 479             line = raw_input(g['decorated_name'])
 480         else:
 481             line = raw_input()
 482         try:
 483             cmd = line.split()[0]
 484         except:
 485             cmd = ''
 486         # Save cmd to global variable and call process
 487         g['stuff'] = ' '.join(line.split()[1:])
 488         process(cmd)()
 489         if cmd in ['switch','t','rt','rep']:
 490             g['prefix'] = False
 491         else:
 492             g['prefix'] = True
 493
 494
 495 def stream(domain, args, name='Rainbow Stream'):
 496     """
 497     Track the stream
 498     """
 499
 500     # The Logo
 501     art_dict = {
 502         USER_DOMAIN: name,
 503         PUBLIC_DOMAIN: args.track_keywords,
 504         SITE_DOMAIN: 'Site Stream',
 505     }
 506     ascii_art(art_dict[domain])
 507
 508     # These arguments are optional:
 509     stream_args = dict(
 510         timeout=args.timeout,
 511         block=not args.no_block,
 512         heartbeat_timeout=args.heartbeat_timeout)
 513
 514     # Track keyword
 515     query_args = dict()
 516     if args.track_keywords:
 517         query_args['track'] = args.track_keywords
 518
 519     # Get stream
 520     stream = TwitterStream(
 521         auth=authen(),
 522         domain=domain,
 523         **stream_args)
 524
 525     if domain == USER_DOMAIN:
 526         tweet_iter = stream.user(**query_args)
 527     elif domain == SITE_DOMAIN:
 528         tweet_iter = stream.site(**query_args)
 529     else:
 530         if args.track_keywords:
 531             tweet_iter = stream.statuses.filter(**query_args)
 532         else:
 533             tweet_iter = stream.statuses.sample()
 534
 535     # Iterate over the stream.
 536     for tweet in tweet_iter:
 537         if tweet is None:
 538             printNicely("-- None --")
 539         elif tweet is Timeout:
 540             printNicely("-- Timeout --")
 541         elif tweet is HeartbeatTimeout:
 542             printNicely("-- Heartbeat Timeout --")
 543         elif tweet is Hangup:
 544             printNicely("-- Hangup --")
 545         elif tweet.get('text'):
 546             draw(t=tweet, keyword=args.track_keywords, fil=args.filter, ig=args.ignore)
 547
 548
 549 def fly():
 550     """
 551     Main function
 552     """
 553     # Spawn stream process
 554     args = parse_arguments()
 555     get_decorated_name()
 556     p = Process(target=stream, args=(USER_DOMAIN, args, g['original_name']))
 557     p.start()
 558
 559     # Start listen process
 560     time.sleep(0.5)
 561     g['reset'] = True
 562     g['prefix'] = True
 563     g['stream_pid'] = p.pid
 564     listen()
 565