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