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