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         print(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         print(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         print(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         print(green('Okay it\'s gone.'))
 241     except:
 242         print(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             print('Newest', SEARCH_MAX_RECORD, 'tweet:')
 254             for i in reversed(xrange(SEARCH_MAX_RECORD)):
 255                 draw(t=rel[i], keyword=g['stuff'].strip()[1:])
 256             printNicely('')
 257         else:
 258             print(red('A keyword should be a hashtag (like \'#AKB48\')'))
 259     except:
 260         print(red('Sorry I can\'t understand.'))
 261
 262
 263 def friend():
 264     """
 265     List of friend (following)
 266     """
 267     t = Twitter(auth=authen())
 268     g['friends'] = t.friends.ids()['ids']
 269     for i in g['friends']:
 270         screen_name = t.users.lookup(user_id=i)[0]['screen_name']
 271         user = cycle_color('@' + screen_name)
 272         print(user, end=' ')
 273     print('\n')
 274
 275
 276 def follower():
 277     """
 278     List of follower
 279     """
 280     t = Twitter(auth=authen())
 281     g['followers'] = t.followers.ids()['ids']
 282     for i in g['followers']:
 283         screen_name = t.users.lookup(user_id=i)[0]['screen_name']
 284         user = cycle_color('@' + screen_name)
 285         print(user, end=' ')
 286     print('\n')
 287
 288
 289 def help():
 290     """
 291     Help
 292     """
 293     usage = '''
 294     Hi boss! I'm ready to serve you right now!
 295     -------------------------------------------------------------
 296     "home" will show your timeline. "home 7" will show 7 tweet.
 297     "view @bob" will show your friend @bob's home.
 298     "t oops" will tweet "oops" immediately.
 299     "rt 12345" will retweet to tweet with id "12345".
 300     "rep 12345 oops" will reply "oops" to tweet with id "12345".
 301     "del 12345" will delete tweet with id "12345".
 302     "s #AKB48" will search for "AKB48" and return 5 newest tweet.
 303     "fr" will list out your following people.
 304     "fl" will list out your followers.
 305     "h" will show this help again.
 306     "c" will clear the terminal.
 307     "q" will exit.
 308     -------------------------------------------------------------
 309     Have fun and hang tight!
 310     '''
 311     printNicely(usage)
 312
 313
 314 def clear():
 315     """
 316     Clear screen
 317     """
 318     os.system('clear')
 319
 320
 321 def quit():
 322     """
 323     Exit all
 324     """
 325     os.system('rm -rf rainbow.db')
 326     os.kill(g['stream_pid'], signal.SIGKILL)
 327     sys.exit()
 328
 329
 330 def reset():
 331     """
 332     Reset prefix of line
 333     """
 334     g['prefix'] = True
 335
 336
 337 def process(cmd):
 338     """
 339     Process switch
 340     """
 341     return dict(zip(
 342         cmdset,
 343         [
 344             home,
 345             view,
 346             tweet,
 347             retweet,
 348             reply,
 349             delete,
 350             search,
 351             friend,
 352             follower,
 353             help,
 354             clear,
 355             quit
 356         ]
 357     )).get(cmd, reset)
 358
 359
 360 def listen():
 361     init_interactive_shell(cmdset)
 362     first = True
 363     while True:
 364         if g['prefix'] and not first:
 365             line = raw_input(g['decorated_name'])
 366         else:
 367             line = raw_input()
 368         try:
 369             cmd = line.split()[0]
 370         except:
 371             cmd = ''
 372         # Save cmd to global variable and call process
 373         g['stuff'] = ' '.join(line.split()[1:])
 374         process(cmd)()
 375         first = False
 376
 377
 378 def stream():
 379     """
 380     Track the stream
 381     """
 382     args = parse_arguments()
 383
 384     # The Logo
 385     ascii_art()
 386     # These arguments are optional:
 387     stream_args = dict(
 388         timeout=args.timeout,
 389         block=not args.no_block,
 390         heartbeat_timeout=args.heartbeat_timeout)
 391
 392     # Track keyword
 393     query_args = dict()
 394     if args.track_keywords:
 395         query_args['track'] = args.track_keywords
 396
 397     # Get stream
 398     stream = TwitterStream(
 399         auth=authen(),
 400         domain=DOMAIN,
 401         **stream_args)
 402     tweet_iter = stream.user(**query_args)
 403
 404     # Iterate over the sample stream.
 405     for tweet in tweet_iter:
 406         if tweet is None:
 407             printNicely("-- None --")
 408         elif tweet is Timeout:
 409             printNicely("-- Timeout --")
 410         elif tweet is HeartbeatTimeout:
 411             printNicely("-- Heartbeat Timeout --")
 412         elif tweet is Hangup:
 413             printNicely("-- Hangup --")
 414         elif tweet.get('text'):
 415             draw(t=tweet)
 416
 417
 418 def fly():
 419     """
 420     Main function
 421     """
 422     get_decorated_name()
 423     g['prefix'] = True
 424     p = Process(target=stream)
 425     p.start()
 426     g['stream_pid'] = p.pid
 427     listen()