diff options
| -rwxr-xr-x | main.py | 83 |
1 files changed, 58 insertions, 25 deletions
@@ -15,6 +15,22 @@ print("root_dir:", root_dir) print("filings_dir:", filings_dir) print("clean_filings_dir:", clean_filings_dir) + +def make_batches(lst, batch_size): + res = [] + batch = [] + for i in range(len(lst)): + batch.append(lst[i]) + if (i + 1) % batch_size == 0: + res.append(batch) + batch = [] + + if len(batch) > 0: + res.append(batch) + + return res + + # Read CSV data into a list of dicts csv_filename = os.path.join(root_dir, "data", "progress.csv") print("Reading CSV data from ", csv_filename) @@ -40,7 +56,6 @@ for row in csv_data_rows: cik = row["cik"] filings = Filing(cik=cik, filing_type="10-k", count=45) - filings.save(filings_dir) input_file_dir = os.path.join(filings_dir, cik, "10-k") @@ -56,6 +71,7 @@ for row in csv_data_rows: data_object = { "ticker": ticker, "cik": cik, + "filings_obj": filings, "input_filename": input_filename, "output_filename": output_filename, "input_string": None, @@ -63,42 +79,59 @@ for row in csv_data_rows: } data_object_lst.append(data_object) -# Read files and store content in data objects -print("Reading filing data into memory") -for data_object in data_object_lst: - input_filename = data_object["input_filename"] - with open(input_filename, "r+") as rawFile: - input_file_string = rawFile.read() - data_object["input_string"] = input_file_string - # Multiprocessing worker function # Stores result in data objects def process(data_object): input_string = data_object["input_string"] - output_string = BeautifulSoup(input_file_string, "lxml").text + output_string = BeautifulSoup(input_string, "lxml").text data_object["output_string"] = output_string return data_object -print("Parsing filing data with lxml") +def process_data_object_batch(batch): + # Fetch filings over HTTP + for data_object in batch: + filings = data_object["filings_obj"] + filings.save(filings_dir) + del data_object["filings_obj"] + + # Read files and store content in data objects + print("Reading filing data into memory") + for data_object in batch: + input_filename = data_object["input_filename"] + with open(input_filename, "r+") as rawFile: + input_file_string = rawFile.read() + data_object["input_string"] = input_file_string + + print("Parsing filing data with lxml") + + # Multiprocessing, 4 cores + with Pool(4) as pool: + result_batch = pool.map(process, batch) + + # No multiprocessing, does the same thing + # result_batch = map(process, batch) + + # Write output to disk from data objects + print("Writing filing data to disk") + for data_object in result_batch: + output_filename = data_object["output_filename"] + output_dir = os.path.dirname(output_filename) + os.makedirs(output_dir, exist_ok=True) + with open(output_filename, "w") as newFile: + output_string = data_object["output_string"] + newFile.write(output_string) + + for data_object in result_batch: + data_object.clear() # Clean up memory -# Multiprocessing, 4 cores -with Pool(4) as pool: - result_data_objects = pool.map(process, data_object_lst) -# No multiprocessing, does the same thing -#result_data_objects = map(process, data_object_lst) +batches = make_batches(data_object_lst, 5) -# Write output to disk from data objects -print("Writing filing data to disk") -for data_object in result_data_objects: - output_filename = data_object["output_filename"] - output_dir = os.path.dirname(output_filename) - os.makedirs(output_dir, exist_ok=True) - with open(output_filename, "w") as newFile: - output_string = data_object["output_string"] - newFile.write(output_string) +for i, batch in enumerate(batches): + print("=== Processing batch #{}".format(i)) + process_data_object_batch(batch) print("Done.") |