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, argparse, sys, time, signal
   9
  10 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
  11 from twitter.api import *
  12 from twitter.oauth import OAuth, read_token_file
  13 from twitter.oauth_dance import oauth_dance
  14 from twitter.util import printNicely
  15 from dateutil import parser
  16
  17 from .colors import *
  18 from .config import *
  19
  20 g = {}
  21
  22 def draw(t):
  23     """
  24     Draw the rainbow
  25     """
  26     # Retrieve tweet
  27     text = t['text']
  28     screen_name = t['user']['screen_name']
  29     name = t['user']['name']
  30     created_at = t['created_at']
  31     date = parser.parse(created_at)
  32     time = date.strftime('%Y/%m/%d %H:%M:%S')
  33
  34     # Format info
  35     user = cycle_color(name) + grey(' ' + '@' + screen_name + ' ')
  36     clock = grey('[' + time + ']')
  37     tweet = text.split()
  38     tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
  39     tweet = map(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
  40     tweet = map(lambda x: cyan(x) if x[0:7] == 'http://' else x, tweet)
  41     tweet = ' '.join(tweet)
  42
  43     # Draw rainbow
  44     line1 = u"{u:>{uw}}:".format(
  45         u=user,
  46         uw=len(user) + 2,
  47     )
  48     line2 = u"{c:>{cw}}".format(
  49         c=clock,
  50         cw=len(clock) + 2,
  51     )
  52     line3 = '  ' + tweet
  53     line4 = '\n'
  54
  55     printNicely(line1)
  56     printNicely(line2)
  57     printNicely(line3)
  58     printNicely(line4)
  59
  60
  61 def parse_arguments():
  62     """
  63     Parse the arguments
  64     """
  65
  66     parser = argparse.ArgumentParser(description=__doc__ or "")
  67
  68     parser.add_argument(
  69         '-to',
  70         '--timeout',
  71         help='Timeout for the stream (seconds).')
  72     parser.add_argument(
  73         '-ht',
  74         '--heartbeat-timeout',
  75         help='Set heartbeat timeout.',
  76         default=90)
  77     parser.add_argument(
  78         '-nb',
  79         '--no-block',
  80         action='store_true',
  81         help='Set stream to non-blocking.')
  82     parser.add_argument(
  83         '-tt',
  84         '--track-keywords',
  85         help='Search the stream for specific text.')
  86     return parser.parse_args()
  87
  88
  89 def authen():
  90     """
  91     authenticate with Twitter OAuth
  92     """
  93     # When using rainbow stream you must authorize.
  94     twitter_credential = os.environ.get(
  95         'HOME',
  96         os.environ.get(
  97             'USERPROFILE',
  98             '')) + os.sep + '.rainbow_oauth'
  99     if not os.path.exists(twitter_credential):
 100         oauth_dance("Rainbow Stream",
 101                     CONSUMER_KEY,
 102                     CONSUMER_SECRET,
 103                     twitter_credential)
 104     oauth_token, oauth_token_secret = read_token_file(twitter_credential)
 105     return OAuth(
 106         oauth_token,
 107         oauth_token_secret,
 108         CONSUMER_KEY,
 109         CONSUMER_SECRET)
 110
 111
 112 def get_decorated_name():
 113     """
 114     Beginning of every line
 115     """
 116     t = Twitter(auth=authen())
 117     name = '@' + t.statuses.user_timeline()[-1]['user']['screen_name']
 118     g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
 119
 120
 121 def tweet():
 122     """
 123     Authen and tweet
 124     """
 125     t = Twitter(auth=authen())
 126     t.statuses.update(status=g['stuff'])
 127
 128
 129 def search():
 130     """
 131     Authen and search
 132     """
 133     t = Twitter(auth=authen())
 134     rel = t.search.tweets(q='#' + g['stuff'])['statuses']
 135     printNicely(grey('**************************************************************************************\n'))
 136     print('Newest ',SEARCH_MAX_RECORD, ' tweet: \n')
 137     for i in xrange(5):
 138         draw(t=rel[i])
 139     printNicely(grey('**************************************************************************************\n'))
 140
 141
 142 def help():
 143     """
 144     Print help
 145     """
 146     usage = '''
 147     Hi boss! I'm ready to serve you right now!
 148     ----------------------------------------------------
 149     "!" at the beginning will start to tweet immediately
 150     "/" at the beginning will search for 5 lastest twwet
 151     "?" will print this help once again
 152     "q" will exit
 153     ----------------------------------------------------
 154     Hvae fun and hang tight!
 155     '''
 156     printNicely(usage)
 157     sys.stdout.write(g['decorated_name'])
 158
 159
 160 def quit():
 161     """
 162     Exit all
 163     """
 164     os.kill(g['stream_pid'], signal.SIGKILL)
 165     sys.exit()
 166
 167
 168 def process(line):
 169     """
 170     Process switch by start of line
 171     """
 172     return {
 173         '!' : tweet,
 174         '/' : search,
 175         '?' : help,
 176         'q' : quit,
 177     }.get(line[0],lambda: sys.stdout.write(g['decorated_name']))
 178
 179
 180 def listen(stdin):
 181     """
 182     Listen to user's input
 183     """
 184     for line in iter(stdin.readline, ''):
 185         # Save cmd to global variable and call process
 186         g['stuff'] = line[1:]
 187         process(line)()
 188     stdin.close()
 189
 190
 191 def stream():
 192     """
 193     Track the stream
 194     """
 195     args = parse_arguments()
 196
 197     # The Logo
 198     ascii_art()
 199     print("Tip: Press ENTER and type a '?' to see what command mode can do for you")
 200     print('\n')
 201     # These arguments are optional:
 202     stream_args = dict(
 203         timeout=args.timeout,
 204         block=not args.no_block,
 205         heartbeat_timeout=args.heartbeat_timeout)
 206
 207     # Track keyword
 208     query_args = dict()
 209     if args.track_keywords:
 210         query_args['track'] = args.track_keywords
 211
 212     # Get stream
 213     stream = TwitterStream(
 214         auth = authen(),
 215         domain = 'userstream.twitter.com',
 216         **stream_args)
 217     tweet_iter = stream.user(**query_args)
 218
 219     # Iterate over the sample stream.
 220     for tweet in tweet_iter:
 221         if tweet is None:
 222             printNicely("-- None --")
 223         elif tweet is Timeout:
 224             printNicely("-- Timeout --")
 225         elif tweet is HeartbeatTimeout:
 226             printNicely("-- Heartbeat Timeout --")
 227         elif tweet is Hangup:
 228             printNicely("-- Hangup --")
 229         elif tweet.get('text'):
 230             draw(t=tweet)
 231
 232
 233 def fly():
 234     """
 235     Main function
 236     """
 237     get_decorated_name()
 238     p = Process(target = stream)
 239     p.start()
 240     g['stream_pid'] = p.pid
 241     listen(sys.stdin)
 242