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