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             # Kill old process
 175             os.kill(g['stream_pid'], signal.SIGKILL)
 176             args = parse_arguments()
 177             args.track_keywords = keyword
 178             # Start new process
 179             p = Process(
 180                 target=stream,
 181                 args=(
 182                     PUBLIC_DOMAIN,
 183                     args))
 184             p.start()
 185             g['stream_pid'] = p.pid
 186
 187         # Personal stream
 188         elif target == 'mine':
 189             # Kill old process
 190             os.kill(g['stream_pid'], signal.SIGKILL)
 191             args = parse_arguments()
 192             # Start new process
 193             p = Process(
 194                 target=stream,
 195                 args=(
 196                     USER_DOMAIN,
 197                     args,
 198                     g['original_name']))
 199             p.start()
 200             g['stream_pid'] = p.pid
 201         printNicely(green('Stream switched.'))
 202     except:
 203         printNicely(red('Sorry I can\'t understand.'))
 204     g['prefix'] = False
 205
 206
 207 def home():
 208     """
 209     Home
 210     """
 211     t = Twitter(auth=authen())
 212     num = HOME_TWEET_NUM
 213     if g['stuff'].isdigit():
 214         num = g['stuff']
 215     for tweet in reversed(t.statuses.home_timeline(count=num)):
 216         draw(t=tweet)
 217     printNicely('')
 218
 219
 220 def view():
 221     """
 222     Friend view
 223     """
 224     t = Twitter(auth=authen())
 225     user = g['stuff'].split()[0]
 226     if user[0] == '@':
 227         try:
 228             num = int(g['stuff'].split()[1])
 229         except:
 230             num = HOME_TWEET_NUM
 231         for tweet in reversed(t.statuses.user_timeline(count=num, screen_name=user[1:])):
 232             draw(t=tweet)
 233         printNicely('')
 234     else:
 235         printNicely(red('A name should begin with a \'@\''))
 236
 237
 238 def tweet():
 239     """
 240     Tweet
 241     """
 242     t = Twitter(auth=authen())
 243     t.statuses.update(status=g['stuff'])
 244     g['prefix'] = False
 245
 246
 247 def retweet():
 248     """
 249     ReTweet
 250     """
 251     t = Twitter(auth=authen())
 252     try:
 253         id = int(g['stuff'].split()[0])
 254         tid = db.rainbow_query(id)[0].tweet_id
 255         t.statuses.retweet(id=tid, include_entities=False, trim_user=True)
 256     except:
 257         printNicely(red('Sorry I can\'t retweet for you.'))
 258     g['prefix'] = False
 259
 260
 261 def reply():
 262     """
 263     Reply
 264     """
 265     t = Twitter(auth=authen())
 266     try:
 267         id = int(g['stuff'].split()[0])
 268         tid = db.rainbow_query(id)[0].tweet_id
 269         user = t.statuses.show(id=tid)['user']['screen_name']
 270         status = ' '.join(g['stuff'].split()[1:])
 271         status = '@' + user + ' ' + status.decode('utf-8')
 272         t.statuses.update(status=status, in_reply_to_status_id=tid)
 273     except:
 274         printNicely(red('Sorry I can\'t understand.'))
 275     g['prefix'] = False
 276
 277
 278 def delete():
 279     """
 280     Delete
 281     """
 282     t = Twitter(auth=authen())
 283     try:
 284         id = int(g['stuff'].split()[0])
 285         tid = db.rainbow_query(id)[0].tweet_id
 286         t.statuses.destroy(id=tid)
 287         printNicely(green('Okay it\'s gone.'))
 288     except:
 289         printNicely(red('Sorry I can\'t delete this tweet for you.'))
 290
 291
 292 def search():
 293     """
 294     Search
 295     """
 296     t = Twitter(auth=authen())
 297     try:
 298         if g['stuff'][0] == '#':
 299             rel = t.search.tweets(q=g['stuff'])['statuses']
 300             if len(rel):
 301                 printNicely('Newest tweets:')
 302                 for i in reversed(xrange(SEARCH_MAX_RECORD)):
 303                     draw(t=rel[i], keyword=g['stuff'].strip()[1:])
 304                 printNicely('')
 305             else:
 306                 printNicely(magenta('I\'m afraid there is no result'))
 307         else:
 308             printNicely(red('A keyword should be a hashtag (like \'#AKB48\')'))
 309     except:
 310         printNicely(red('Sorry I can\'t understand.'))
 311
 312
 313 def friend():
 314     """
 315     List of friend (following)
 316     """
 317     t = Twitter(auth=authen())
 318     g['friends'] = t.friends.ids()['ids']
 319     for i in g['friends']:
 320         screen_name = t.users.lookup(user_id=i)[0]['screen_name']
 321         user = cycle_color('@' + screen_name)
 322         print(user, end=' ')
 323     printNicely('')
 324
 325
 326 def follower():
 327     """
 328     List of follower
 329     """
 330     t = Twitter(auth=authen())
 331     g['followers'] = t.followers.ids()['ids']
 332     for i in g['followers']:
 333         screen_name = t.users.lookup(user_id=i)[0]['screen_name']
 334         user = cycle_color('@' + screen_name)
 335         print(user, end=' ')
 336     printNicely('')
 337
 338
 339 def help():
 340     """
 341     Help
 342     """
 343     usage = '''
 344     Hi boss! I'm ready to serve you right now!
 345     -------------------------------------------------------------
 346     You are already on your personal stream:
 347       "switch public #AKB" will switch to public stream and follow "AKB" keyword.
 348       "switch mine" will switch back to your personal stream.
 349     For more action:
 350       "home" will show your timeline. "home 7" will show 7 tweet.
 351       "view @bob" will show your friend @bob's home.
 352       "t oops" will tweet "oops" immediately.
 353       "rt 12345" will retweet to tweet with id "12345".
 354       "rep 12345 oops" will reply "oops" to tweet with id "12345".
 355       "del 12345" will delete tweet with id "12345".
 356       "s #AKB48" will search for "AKB48" and return 5 newest tweet.
 357       "fr" will list out your following people.
 358       "fl" will list out your followers.
 359       "h" will show this help again.
 360       "c" will clear the terminal.
 361       "q" will exit.
 362     -------------------------------------------------------------
 363     Have fun and hang tight!
 364     '''
 365     printNicely(usage)
 366
 367
 368 def clear():
 369     """
 370     Clear screen
 371     """
 372     os.system('clear')
 373     g['prefix'] = False
 374
 375
 376 def quit():
 377     """
 378     Exit all
 379     """
 380     os.system('rm -rf rainbow.db')
 381     os.kill(g['stream_pid'], signal.SIGKILL)
 382     sys.exit()
 383
 384
 385 def reset():
 386     """
 387     Reset prefix of line
 388     """
 389     if g['reset']:
 390         printNicely(green('Need tips ? Type "h" and hit Enter key!'))
 391     g['prefix'] = True
 392     g['reset'] = False
 393
 394
 395 def process(cmd):
 396     """
 397     Process switch
 398     """
 399     return dict(zip(
 400         cmdset,
 401         [
 402             switch,
 403             home,
 404             view,
 405             tweet,
 406             retweet,
 407             reply,
 408             delete,
 409             search,
 410             friend,
 411             follower,
 412             help,
 413             clear,
 414             quit
 415         ]
 416     )).get(cmd, reset)
 417
 418
 419 def listen():
 420     """
 421     Listen to user's input
 422     """
 423     init_interactive_shell(cmdset)
 424     first = True
 425     while True:
 426         if g['prefix'] and not first:
 427             line = raw_input(g['decorated_name'])
 428         else:
 429             line = raw_input()
 430         try:
 431             cmd = line.split()[0]
 432         except:
 433             cmd = ''
 434         # Save cmd to global variable and call process
 435         g['stuff'] = ' '.join(line.split()[1:])
 436         process(cmd)()
 437         first = False
 438
 439
 440 def stream(domain, args, name='Rainbow Stream'):
 441     """
 442     Track the stream
 443     """
 444     # The Logo
 445     art_dict = {
 446         USER_DOMAIN: name,
 447         PUBLIC_DOMAIN: args.track_keywords,
 448         SITE_DOMAIN: 'Site Stream',
 449     }
 450     ascii_art(art_dict[domain])
 451     # These arguments are optional:
 452     stream_args = dict(
 453         timeout=args.timeout,
 454         block=not args.no_block,
 455         heartbeat_timeout=args.heartbeat_timeout)
 456
 457     # Track keyword
 458     query_args = dict()
 459     if args.track_keywords:
 460         query_args['track'] = args.track_keywords
 461
 462     # Get stream
 463     stream = TwitterStream(
 464         auth=authen(),
 465         domain=domain,
 466         **stream_args)
 467
 468     if domain == USER_DOMAIN:
 469         tweet_iter = stream.user(**query_args)
 470     elif domain == SITE_DOMAIN:
 471         tweet_iter = stream.site(**query_args)
 472     else:
 473         if args.track_keywords:
 474             tweet_iter = stream.statuses.filter(**query_args)
 475         else:
 476             tweet_iter = stream.statuses.sample()
 477
 478     # Iterate over the stream.
 479     for tweet in tweet_iter:
 480         if tweet is None:
 481             printNicely("-- None --")
 482         elif tweet is Timeout:
 483             printNicely("-- Timeout --")
 484         elif tweet is HeartbeatTimeout:
 485             printNicely("-- Heartbeat Timeout --")
 486         elif tweet is Hangup:
 487             printNicely("-- Hangup --")
 488         elif tweet.get('text'):
 489             draw(t=tweet, keyword=args.track_keywords)
 490
 491
 492 def fly():
 493     """
 494     Main function
 495     """
 496     # Spawn stream process
 497     args = parse_arguments()
 498     get_decorated_name()
 499     p = Process(target=stream, args=(USER_DOMAIN, args, g['original_name']))
 500     p.start()
 501
 502     # Start listen process
 503     g['prefix'] = True
 504     g['reset'] = True
 505     g['stream_pid'] = p.pid
 506     listen()
 507
 508
 509 def multiplexer():
 510     """
 511     Multiplexer
 512     """
 513     os.system('screen -S rainbow fly')