diff --git a/atdb/communication.py b/atdb/communication.py index b466c2027a112fa859ec8e33de51f53b34c04bc4..68f8d8b02978622f40d2874a7e427b9f87b2a081 100644 --- a/atdb/communication.py +++ b/atdb/communication.py @@ -127,8 +127,9 @@ class APIConnector: """ Change the whole task content """ - return self._request_path("PUT", f"tasks/{task_id}", - content={"size_processed": processed_size}) + return self._request_path( + "PUT", f"tasks/{task_id}", content={"size_processed": processed_size} + ) def change_task_status(self, task_id, status) -> None: """ diff --git a/atdb/fix.py b/atdb/fix.py index 196a00d7b38534e54ad1442a9a762f6d3957a617..efff84aa5293d36fbe0ab22df950e892b2e6416f 100644 --- a/atdb/fix.py +++ b/atdb/fix.py @@ -1,44 +1,71 @@ -import atdb.communication as com +""" +Fix command module +""" import logging +import atdb.communication as com + + +logger = logging.getLogger("fix") + -logger = logging.getLogger('fix') def aggregate_on_tree(tree, field): + """ + Aggregated values with a given field name from a dict tree + """ if isinstance(tree, dict) and field in tree: return tree[field] - elif isinstance(tree, dict): + if isinstance(tree, dict): total = 0 for value in tree.values(): total += aggregate_on_tree(value, field) return total - elif isinstance(tree, list): + if isinstance(tree, list): total = 0 for item in tree: total += aggregate_on_tree(item, field) return total - else: - return 0 + + return 0 + def compute_output_sizes(outputs): + """ + Computes the size of the output files + """ if outputs is not None: - return aggregate_on_tree({key: value for key, value in outputs.items() if key != 'ingest'}, 'size') - else: - return 0 + return aggregate_on_tree( + {key: value for key, value in outputs.items() if key != "ingest"}, "size" + ) + return 0 + def fix_computed_sizes(connector, dry_run=True): - for task in connector.list_iter('tasks'): - task_id = task['id'] - size_before = task['size_processed'] - total_output_size = compute_output_sizes(task['outputs']) - task['size_processed'] = total_output_size + """ + Fix the size of the computed task + """ + for task in connector.list_iter("tasks"): + task_id = task["id"] + size_before = task["size_processed"] + total_output_size = compute_output_sizes(task["outputs"]) + task["size_processed"] = total_output_size if not dry_run: if size_before != total_output_size: connector.update_task_processed_size(task_id, total_output_size) else: if size_before != total_output_size: - logger.info('Dry run: Size updated for %s from %s to %s', task_id, - size_before, total_output_size) + logger.info( + "Dry run: Size updated for %s from %s to %s", + task_id, + size_before, + total_output_size, + ) def fix(args): + """ + Fix command + + Changes task fields to be consistent with each others + """ connector = com.APIConnector.from_args(args) - fix_computed_sizes(connector, dry_run=args.dry_run) \ No newline at end of file + fix_computed_sizes(connector, dry_run=args.dry_run)