fix python2.6 compat and refacto per as @Adonoho's #201

[z_archive/twitter.git] / twitter / stream_example.py

diff --git a/twitter/stream_example.py b/twitter/stream_example.py
index dd4fffe3083753d1976a230474220e0c94fd2e66..2888aced4c05e8681ab5bca63c1dcc640e991699 100644 (file)
--- a/twitter/stream_example.py
+++ b/twitter/stream_example.py
@@ -4,30 +4,57 @@ from the "sample" stream as fast as possible.
 
 USAGE
 
-  twitter-stream-example <username> <password>
+  stream-example -t <token> -ts <token_secret> -ck <consumer_key> -cs <consumer_secret>
 
 """
 
 from __future__ import print_function
 
-import sys
+import argparse
 
-from .stream import TwitterStream
-from .auth import UserPassAuth
-from .util import printNicely
+from twitter.stream import TwitterStream
+from twitter.oauth import OAuth
+from twitter.util import printNicely
 
-def main(args=sys.argv[1:]):
-    if not args[1:]:
-        print(__doc__)
-        return 1
 
-    # When using twitter stream you must authorize. UserPass or OAuth.
-    stream = TwitterStream(auth=UserPassAuth(args[0], args[1]))
+def parse_arguments():
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('-t',  '--token', help='The Twitter Access Token.')
+    parser.add_argument('-ts', '--token_secret', help='The Twitter Access Token Secret.')
+    parser.add_argument('-ck', '--consumer_key', help='The Twitter Consumer Key.')
+    parser.add_argument('-cs', '--consumer_secret', help='The Twitter Consumer Secret.')
+    parser.add_argument('-us', '--user_stream', action='store_true', help='Connect to the user stream endpoint.')
+    parser.add_argument('-ss', '--site_stream', action='store_true', help='Connect to the site stream endpoint.')
+
+    return parser.parse_args()
+
+def main():
+    args = parse_arguments()
+
+    if not all((args.token, args.token_secret, args.consumer_key, args.consumer_secret)):
+        print(__doc__)
+        return 2
+
+    # When using twitter stream you must authorize.
+    auth = OAuth(args.token, args.token_secret, args.consumer_key, args.consumer_secret)
+    if args.user_stream:
+        stream = TwitterStream(auth=auth, domain='userstream.twitter.com')
+        tweet_iter = stream.user()
+    elif args.site_stream:
+        stream = TwitterStream(auth=auth, domain='sitestream.twitter.com')
+        tweet_iter = stream.site()
+    else:
+        stream = TwitterStream(auth=auth, timeout=60.0)
+        tweet_iter = stream.statuses.sample()
 
     # Iterate over the sample stream.
-    tweet_iter = stream.statuses.sample()
     for tweet in tweet_iter:
         # You must test that your tweet has text. It might be a delete
         # or data message.
         if tweet.get('text'):
             printNicely(tweet['text'])
+
+if __name__ == '__main__':
+    main()