extractor.py

#!/usr/bin/env python3
from datetime import datetime
import json
import os
import re
import argparse
import csv
import copy
import sys
import gzip

strptime = datetime.strptime

class attriObject:
    """Class object for attribute parser."""
    def __init__(self, string):
        self.value = re.split(":", string)
        self.title = self.value[-1]

    def getElement(self, json_object):
        found = [json_object]
        for entry in self.value:
            for index in range(len(found)):
                try:
                    found[index] = found[index][entry]
                except (TypeError, KeyError):
                    print("'{0}' is not a valid json entry.".format(":".join(self.value)))
                    sys.exit()

                #If single search object is a list, search entire list. Error if nested lists.
                if isinstance(found[index], list):
                    if len(found) > 1:
                        raise Exception("Extractor currently does not handle nested lists.")
                    found = found[index]

        return found

def tweets_files(string, path):
    """Iterates over json files in path."""
    for filename in os.listdir(path):
        if re.match(string, filename) and ".json" in filename:
            f = gzip.open if ".gz" in filename else open
            yield path + filename, f
            
            Ellipsis

def parse(args):
    with open(args.output, 'w+', encoding="utf-8") as output:
        csv_writer = csv.writer(output, dialect=args.dialect)
        csv_writer.writerow([a.title for a in args.attributes])
        count = 0
        tweets = set()

        for filename, f in tweets_files(args.string, args.path):
            print("parsing", filename)
            with f(filename, 'rb') as data_file:
                for line in data_file:
                    try:
                        json_object = json.loads(line.decode("utf-8"))
                    except ValueError:
                        print("Error in", filename, "entry incomplete.")
                        continue

                    #Check for duplicates
                    identity = json_object['id']
                    if identity in tweets:
                        continue
                    tweets.add(identity)

                    #Check for time restrictions.
                    if args.start or args.end:
                        tweet_time = strptime(json_object['created_at'],'%a %b %d %H:%M:%S +0000 %Y')
                        if args.start and args.start > tweet_time:
                            continue
                        if args.end and args.end < tweet_time:
                            continue

                    #Check for hashtag.
                    if args.hashtag:
                        for entity in json_object['entities']["hashtags"]:
                            if entity['text'].lower() == args.hashtag:
                                break
                        else:
                            continue

                    count += extract(json_object, args, csv_writer)

        print("Searched", len(tweets), "tweets and recorded", count, "items.")
        print("largest id:", max(tweets))

def extract(json_object, args, csv_writer):
    """Extract and write found attributes."""
    found = [[]]
    for attribute in args.attributes:
        item = attribute.getElement(json_object)
        if len(item) == 0:
            for row in found:
                row.append("NA")
        else:
            found1 = []
            for value in item:
                if value is None:
                    value = "NA"
                new = copy.deepcopy(found)
                for row in new:
                    row.append(value)
                found1.extend(new)
            found = found1    

    for row in found:
        
        csv_writer.writerow(row)
    return len(found)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Extracts attributes from tweets.')
    parser.add_argument("attributes", nargs='*', help="Attributes to search for. Attributes inside nested inside other attributes should be seperated by a colon. Example: user:screen_name, entities:hashtags:text.")
    parser.add_argument("-dialect", default="excel", help="Sets dialect for csv output. Defaults to excel. See python module csv.list_dialects()")
    parser.add_argument("-string", default="", help="Regular expression for files to parse. Defaults to empty string.")
    parser.add_argument("-path", default="./", help="Optional path to folder containing tweets. Defaults to current folder.")
    parser.add_argument("-output", default="output.csv", help="Optional file to output results. Defaults to output.csv.")
    parser.add_argument("-start", default="", help="Define start date for tweets. Format (mm:dd:yyyy)")
    parser.add_argument("-end", default="", help="Define end date for tweets. Format (mm:dd:yyyy)")
    parser.add_argument("-hashtag", default="", help="Define a hashtag that must be in parsed tweets.")
    args = parser.parse_args()

    if not args.path.endswith("/"):
        args.path += "/"

    args.start = strptime(args.start, '%m:%d:%Y') if args.start else False
    args.end = strptime(args.end, '%m:%d:%Y') if args.end else False
    args.attributes = [attriObject(i) for i in args.attributes]
    args.string = re.compile(args.string)
    args.hashtag = args.hashtag.lower()
    
    parse(args)