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