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