rainbowstream/rainbow.py

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