Refactor filter_date util, add support for max date

milesmcc · milesmcc · commit 915e2958e892 · 2020-10-19T19:55:21.000Z
diff --git a/utils/filter_date.py b/utils/filter_date.py
@@ -1,49 +1,46 @@
 #!/usr/bin/env python
 """
-Given a minimum date, filter out all tweets after this date.
+Given a minimum and/or maximum date, filter out all tweets after this date.
 
 For example, if a hashtag was used for another event before the one you're
 interested in, you can filter out the old ones.
 
 Example usage:
-utils\filter_date.py --mindate 1-may-2014 tweets.jsonl > filtered.jsonl
+utils/filter_date.py --mindate 1-may-2014 tweets.jsonl > filtered.jsonl
 """
 from __future__ import print_function
 
 import sys
 import json
 import fileinput
-import dateutil.parser
+import argparse
+import datetime
+from dateutil.parser import parse
 
 
-# parse command-line args
-mindate = dateutil.parser.parse("1-January-2012")
-# if args include --mindate, get mindate and remove first two args,
-# leaving file name(s) (if any) in args
-if len(sys.argv) > 1:
-    if sys.argv[1] == "--mindate":
-        mindate = dateutil.parser.parse(sys.argv[2])
-        del sys.argv[0]
-        del sys.argv[0]
+def filter_input(mindate, maxdate, files):
+    mindate = parse(mindate) if mindate is not None else datetime.datetime.min
+    maxdate = parse(maxdate) if maxdate is not None else datetime.datetime.max
 
-# fh = open('date_filtered.jsonl', 'w')
-# kept, discarded = 0, 0
+    for line in fileinput.input(files):
+        tweet = json.loads(line)
 
-for line in fileinput.input():
-    tweet = json.loads(line)
+        created_at = parse(tweet["created_at"])
+        created_at = created_at.replace(tzinfo=None)
 
-    created_at = dateutil.parser.parse(tweet["created_at"])
-    created_at = created_at.replace(tzinfo=None)
+        if mindate < created_at and maxdate > created_at:
+            print(json.dumps(tweet))
 
-    # print(created_at, mindate, created_at >= mindate)
-    if (created_at >= mindate):
-        print(json.dumps(tweet))
-        # fh.write(json.dumps(tweet))
-        # fh.write("\n")
-        # kept += 1
-    # else:
-        # discarded += 1
 
-# print("Kept", kept, "tweets and discarded", discarded)
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--mindate", help="the minimum date", default=None)
+    parser.add_argument("--maxdate", help="the maximum date", default=None)
+    parser.add_argument("files", nargs="?", default=[])
+    args = parser.parse_args()
 
-# End of file
+    filter_input(args.mindate, args.maxdate, args.files)
+
+
+if __name__ == "__main__":
+    main()
