add cmdline tool

[rainbowstream.git] / rainbowstream / rainbow.py

diff --git a/rainbowstream/rainbow.py b/rainbowstream/rainbow.py
index 2a070e35814a8d517e943e058e6c463868958737..f489c3f67a4a3b097baf3e4f2f4ebd1ebe21c3c5 100644 (file)
--- a/rainbowstream/rainbow.py
+++ b/rainbowstream/rainbow.py
@@ -3,12 +3,12 @@ Colorful user's timeline stream
 """
 
 from __future__ import print_function
+from multiprocessing import Process
 
-import os
-import os.path
-import argparse
+import os, os.path, argparse, sys, time, signal
 
 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
+from twitter.api import *
 from twitter.oauth import OAuth, read_token_file
 from twitter.oauth_dance import oauth_dance
 from twitter.util import printNicely
@@ -17,6 +17,7 @@ from dateutil import parser
 from .colors import *
 from .config import *
 
+g = {}
 
 def draw(t):
     """
@@ -31,7 +32,7 @@ def draw(t):
     time = date.strftime('%Y/%m/%d %H:%M:%S')
 
     # Format info
-    user = cycle_color(name + ' ' + '@' + screen_name + ' ')
+    user = cycle_color(name) + grey(' ' + '@' + screen_name + ' ')
     clock = grey('[' + time + ']')
     tweet = text.split()
     tweet = map(lambda x: grey(x) if x == 'RT' else x, tweet)
@@ -51,7 +52,10 @@ def draw(t):
     line3 = '  ' + tweet
     line4 = '\n'
 
-    return line1, line2, line3, line4
+    printNicely(line1)
+    printNicely(line2)
+    printNicely(line3)
+    printNicely(line4)
 
 
 def parse_arguments():
@@ -82,12 +86,10 @@ def parse_arguments():
     return parser.parse_args()
 
 
-def stream():
-    args = parse_arguments()
-
-    # The Logo
-    ascii_art()
-
+def authen():
+    """
+    authenticate with Twitter OAuth
+    """
     # When using rainbow stream you must authorize.
     twitter_credential = os.environ.get(
         'HOME',
@@ -100,12 +102,102 @@ def stream():
                     CONSUMER_SECRET,
                     twitter_credential)
     oauth_token, oauth_token_secret = read_token_file(twitter_credential)
-    auth = OAuth(
+    return OAuth(
         oauth_token,
         oauth_token_secret,
         CONSUMER_KEY,
         CONSUMER_SECRET)
 
+
+def get_decorated_name():
+    """
+    Beginning of every line
+    """
+    t = Twitter(auth=authen())
+    name = '@' + t.statuses.user_timeline()[-1]['user']['screen_name']
+    g['decorated_name'] = grey('[') + grey(name) + grey(']: ')
+
+
+def tweet():
+    """
+    Authen and tweet
+    """
+    t = Twitter(auth=authen())
+    t.statuses.update(status=g['stuff'])
+
+
+def search():
+    """
+    Authen and search
+    """
+    t = Twitter(auth=authen())
+    rel = t.search.tweets(q='#' + g['stuff'])['statuses']
+    printNicely(grey('**************************************************************************************\n'))
+    print('Newest ',SEARCH_MAX_RECORD, ' tweet: \n')
+    for i in xrange(5):
+        draw(t=rel[i])
+    printNicely(grey('**************************************************************************************\n'))
+
+
+def help():
+    """
+    Print help
+    """
+    usage = '''
+    Hi boss! I'm ready to serve you right now!
+    ----------------------------------------------------
+    "!" at the beginning will start to tweet immediately
+    "/" at the beginning will search for 5 lastest twwet
+    "?" will print this help once again
+    "q" will exit
+    ----------------------------------------------------
+    Hvae fun and hang tight!
+    '''
+    printNicely(usage)
+    sys.stdout.write(g['decorated_name'])
+
+
+def quit():
+    """
+    Exit all
+    """
+    os.kill(g['stream_pid'], signal.SIGKILL)
+    sys.exit()
+
+
+def process(line):
+    """
+    Process switch by start of line
+    """
+    return {
+        '!' : tweet,
+        '/' : search,
+        '?' : help,
+        'q' : quit,
+    }.get(line[0],lambda: sys.stdout.write(g['decorated_name']))
+
+
+def listen(stdin):
+    """
+    Listen to user's input
+    """
+    for line in iter(stdin.readline, ''):
+        # Save cmd to global variable and call process
+        g['stuff'] = line[1:]
+        process(line)()
+    stdin.close()
+
+
+def stream():
+    """
+    Track the stream
+    """
+    args = parse_arguments()
+
+    # The Logo
+    ascii_art()
+    print("Tip: Press ENTER and type a '?' to see what command mode can do for you")
+    print('\n')
     # These arguments are optional:
     stream_args = dict(
         timeout=args.timeout,
@@ -119,8 +211,8 @@ def stream():
 
     # Get stream
     stream = TwitterStream(
-        auth=auth,
-        domain='userstream.twitter.com',
+        auth = authen(),
+        domain = 'userstream.twitter.com',
         **stream_args)
     tweet_iter = stream.user(**query_args)
 
@@ -135,8 +227,16 @@ def stream():
         elif tweet is Hangup:
             printNicely("-- Hangup --")
         elif tweet.get('text'):
-            line1, line2, line3, line4 = draw(t=tweet)
-            printNicely(line1)
-            printNicely(line2)
-            printNicely(line3)
-            printNicely(line4)
+            draw(t=tweet)
+
+
+def fly():
+    """
+    Main function
+    """
+    get_decorated_name()
+    p = Process(target = stream)
+    p.start()
+    g['stream_pid'] = p.pid
+    listen(sys.stdin)
+