diff --git a/getmentions.py b/getmentions.py index cade951..306363b 100755 --- a/getmentions.py +++ b/getmentions.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 import dbtools - -import operator, re, sys +import argparse, operator, re, sys def getTweets(mode = "@", db = dbtools.dbObject()): handles = dict() @@ -24,20 +23,18 @@ def getTweets(mode = "@", db = dbtools.dbObject()): return handles if __name__ == "__main__": - mode = "@" - if len(sys.argv) > 1: - if len(sys.argv) > 3 or (len(sys.argv) == 3 and "-h" not in sys.argv): - raise ValueError("Invalid arguments passed.") + parser = argparse.ArgumentParser() + g = parser.add_mutually_exclusive_group() + g.add_argument("-t", "--hashtags", help="count only #hashtags", action="store_true") + g.add_argument("-w", "--words", help="count all words", action="store_true") + g.add_argument("-m", "--mentions", help="count only @mentions (default)", action="store_true") + args = parser.parse_args() + if args.hashtags: + mode = "#" + elif args.words: + mode = "" + else: + mode = "@" - path = None - - for arg in sys.argv[1:]: - if arg == "-h": - mode = "#" - if arg == "-w": - mode = "" - else: - path = arg - - for handle, tweets in sorted(list(getTweets(mode,dbtools.dbObject(path=path)).items()), key=operator.itemgetter(1), reverse=True): + for handle, tweets in sorted(list(getTweets(mode,dbtools.dbObject()).items()), key=operator.itemgetter(1), reverse=True): print(handle + "," + str(tweets))