Add batchesmaster

1 files changed, 58 insertions, 25 deletions

diff --git a/main.py b/main.py
index d7938f5..ac31d8a 100755
--- a/main.py
+++ b/main.py
@@ -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.")