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     print('before search')
  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.statuses.user_timeline()[-1]['user']['screen_name']
 158     g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
 159
 160
 161 def home():
 162     """
 163     Home
 164     """
 165     t = Twitter(auth=authen())
 166     num = HOME_TWEET_NUM
 167     if g['stuff'].isdigit():
 168         num = g['stuff']
 169     for tweet in reversed(t.statuses.home_timeline(count=num)):
 170         draw(t=tweet)
 171     printNicely('')
 172
 173
 174 def view():
 175     """
 176     Friend view
 177     """
 178     t = Twitter(auth=authen())
 179     user = g['stuff'].split()[0]
 180     if user[0] == '@':
 181         try:
 182             num = int(g['stuff'].split()[1])
 183         except:
 184             num = HOME_TWEET_NUM
 185         for tweet in reversed(t.statuses.user_timeline(count=num, screen_name=user[1:])):
 186             draw(t=tweet)
 187         printNicely('')
 188     else:
 189         print(red('A name should begin with a \'@\''))
 190
 191
 192 def tweet():
 193     """
 194     Tweet
 195     """
 196     t = Twitter(auth=authen())
 197     t.statuses.update(status=g['stuff'])
 198     g['prefix'] = False
 199
 200
 201 def retweet():
 202     """
 203     ReTweet
 204     """
 205     t = Twitter(auth=authen())
 206     try:
 207         id = int(g['stuff'].split()[0])
 208         tid = db.rainbow_query(id)[0].tweet_id
 209         t.statuses.retweet(id=tid, include_entities=False, trim_user=True)
 210     except:
 211         print(red('Sorry I can\'t retweet for you.'))
 212     g['prefix'] = False
 213
 214
 215 def reply():
 216     """
 217     Reply
 218     """
 219     t = Twitter(auth=authen())
 220     try:
 221         id = int(g['stuff'].split()[0])
 222         tid = db.rainbow_query(id)[0].tweet_id
 223         user = t.statuses.show(id=tid)['user']['screen_name']
 224         status = ' '.join(g['stuff'].split()[1:])
 225         status = '@' + user + ' ' + status.decode('utf-8')
 226         t.statuses.update(status=status, in_reply_to_status_id=tid)
 227     except:
 228         print(red('Sorry I can\'t understand.'))
 229     g['prefix'] = False
 230
 231
 232 def delete():
 233     """
 234     Delete
 235     """
 236     t = Twitter(auth=authen())
 237     try:
 238         id = int(g['stuff'].split()[0])
 239         tid = db.rainbow_query(id)[0].tweet_id
 240         t.statuses.destroy(id=tid)
 241         print(green('Okay it\'s gone.'))
 242     except:
 243         print(red('Sorry I can\'t delete this tweet for you.'))
 244
 245
 246 def search():
 247     """
 248     Search
 249     """
 250     t = Twitter(auth=authen())
 251     try:
 252         if g['stuff'][0] == '#':
 253             rel = t.search.tweets(q=g['stuff'])['statuses']
 254             print('Newest', SEARCH_MAX_RECORD, 'tweet:')
 255             for i in xrange(5):
 256                 draw(t=rel[i], keyword=g['stuff'].strip()[1:])
 257             printNicely('')
 258         else:
 259             print(red('A keyword should be a hashtag (like \'#AKB48\')'))
 260     except:
 261         print(red('Sorry I can\'t understand.'))
 262
 263
 264 def friend():
 265     """
 266     List of friend (following)
 267     """
 268     t = Twitter(auth=authen())
 269     g['friends'] = t.friends.ids()['ids']
 270     for i in g['friends']:
 271         screen_name = t.users.lookup(user_id=i)[0]['screen_name']
 272         user = cycle_color('@' + screen_name)
 273         print(user, end=' ')
 274     print('\n')
 275
 276
 277 def follower():
 278     """
 279     List of follower
 280     """
 281     t = Twitter(auth=authen())
 282     g['followers'] = t.followers.ids()['ids']
 283     for i in g['followers']:
 284         screen_name = t.users.lookup(user_id=i)[0]['screen_name']
 285         user = cycle_color('@' + screen_name)
 286         print(user, end=' ')
 287     print('\n')
 288
 289
 290 def help():
 291     """
 292     Help
 293     """
 294     usage = '''
 295     Hi boss! I'm ready to serve you right now!
 296     -------------------------------------------------------------
 297     "home" will show your timeline. "home 7" will show 7 tweet.
 298     "view @bob" will show your friend @bob's home.
 299     "t oops" will tweet "oops" immediately.
 300     "rt 12345" will retweet to tweet with id "12345".
 301     "rep 12345 oops" will reply "oops" to tweet with id "12345".
 302     "del 12345" will delete tweet with id "12345".
 303     "s #AKB48" will search for "AKB48" and return 5 newest tweet.
 304     "fr" will list out your following people.
 305     "fl" will list out your followers.
 306     "h" will show this help again.
 307     "c" will clear the terminal.
 308     "q" will exit.
 309     -------------------------------------------------------------
 310     Have fun and hang tight!
 311     '''
 312     printNicely(usage)
 313
 314
 315 def clear():
 316     """
 317     Clear screen
 318     """
 319     os.system('clear')
 320
 321
 322 def quit():
 323     """
 324     Exit all
 325     """
 326     db.truncate()
 327     os.kill(g['stream_pid'], signal.SIGKILL)
 328     sys.exit()
 329
 330
 331 def reset():
 332     """
 333     Reset prefix of line
 334     """
 335     g['prefix'] = True
 336
 337
 338 def process(cmd):
 339     """
 340     Process switch
 341     """
 342     return dict(zip(
 343         cmdset,
 344         [
 345             home,
 346             view,
 347             tweet,
 348             retweet,
 349             reply,
 350             delete,
 351             search,
 352             friend,
 353             follower,
 354             help,
 355             clear,
 356             quit
 357         ]
 358     )).get(cmd, reset)
 359
 360
 361 def listen():
 362     init_interactive_shell(cmdset)
 363     first = True
 364     while True:
 365         if g['prefix'] and not first:
 366             line = raw_input(g['decorated_name'])
 367         else:
 368             line = raw_input()
 369         try:
 370             cmd = line.split()[0]
 371         except:
 372             cmd = ''
 373         # Save cmd to global variable and call process
 374         g['stuff'] = ' '.join(line.split()[1:])
 375         process(cmd)()
 376         first = False
 377
 378
 379 def stream():
 380     """
 381     Track the stream
 382     """
 383     args = parse_arguments()
 384
 385     # The Logo
 386     ascii_art()
 387     # These arguments are optional:
 388     stream_args = dict(
 389         timeout=args.timeout,
 390         block=not args.no_block,
 391         heartbeat_timeout=args.heartbeat_timeout)
 392
 393     # Track keyword
 394     query_args = dict()
 395     if args.track_keywords:
 396         query_args['track'] = args.track_keywords
 397
 398     # Get stream
 399     stream = TwitterStream(
 400         auth=authen(),
 401         domain=DOMAIN,
 402         **stream_args)
 403     tweet_iter = stream.user(**query_args)
 404
 405     # Iterate over the sample stream.
 406     for tweet in tweet_iter:
 407         if tweet is None:
 408             printNicely("-- None --")
 409         elif tweet is Timeout:
 410             printNicely("-- Timeout --")
 411         elif tweet is HeartbeatTimeout:
 412             printNicely("-- Heartbeat Timeout --")
 413         elif tweet is Hangup:
 414             printNicely("-- Hangup --")
 415         elif tweet.get('text'):
 416             draw(t=tweet)
 417
 418
 419 def fly():
 420     """
 421     Main function
 422     """
 423     get_decorated_name()
 424     g['prefix'] = True
 425     p = Process(target=stream)
 426     p.start()
 427     g['stream_pid'] = p.pid
 428     listen()