Create batches.py

Author: raphaeljafriLB

labelpandas/batches.py

@@ -0,0 +1,30 @@
+import pandas as pd
+
+def create_batches_dict(table: pandas.core.frame.DataFrame, table_dict:dict, 
+                        global_key_col:str, project_id_col:str, 
+                        project_id:str, global_key_to_data_row_id:dict):
+    """ From a Pandas DataFrame, creates a dictionary where {key=project_id : value=list_of_data_row_ids}
+    Args:
+        table                       :   Required (pandas.core.frame.DataFrame) - Pandas DataFrame                
+        table_dict                  :   Required (dict) - Pandas DataFrame as dict with df.to_dict("records")
+        global_key_col              :   Required (str) - Column name containing the data row global key - defaults to row data
+        project_id_col              :   Required (str) - Column name containing the project ID to batch a given row to
+        project_id                  :   Required (str) - Labelbox project ID to add data rows to - only necessary if no "project_id" column exists
+        global_key_to_data_row_id   :   Required (dict) - Dictionary where {key=global_key : value=data_row_id}
+    Returns:
+        Dictionary where {key=project_id : value=list_of_data_row_ids}
+    """
+    if project_id:
+        project_id_to_batch_dict = {project_id : []}
+    else:
+        project_ids = labelpandas.connector.get_unique_values_function(table=table)
+        project_id_to_batch_dict = {id : [] for id in project_ids}
+    errors = []
+    try:
+        for row in table_dict:
+            id = project_id if project_id else row[project_id_col]
+            data_row_id = global_key_to_data_row_id[row[global_key_col]]
+            project_id_to_batch_dict[id].append(data_row_id)
+    except Exception as e:
+        errors = e
+    return project_id_to_batch_dict, errors