From 210c352ae8a601833a6e34b13972145abc1d49a6 Mon Sep 17 00:00:00 2001 From: sundar-mudupalli-work <106282841+sundar-mudupalli-work@users.noreply.github.com> Date: Wed, 22 May 2024 11:25:56 -0700 Subject: [PATCH] fix: Update to support up to 10K partitions (#1139) * Updated to support 10K partitions * Updated after testing * Fixing syntax issue * Update data_validation/cli_tools.py Co-authored-by: Neha Nene * Update data_validation/cli_tools.py Co-authored-by: Neha Nene * Update data_validation/cli_tools.py Co-authored-by: Neha Nene * Simpler check for value of partition number * Correct check for value of partition number --------- Co-authored-by: Neha Nene --- data_validation/cli_tools.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/data_validation/cli_tools.py b/data_validation/cli_tools.py index b97a09b4..97ff38bc 100644 --- a/data_validation/cli_tools.py +++ b/data_validation/cli_tools.py @@ -954,6 +954,17 @@ def _add_common_arguments(optional_arguments, required_arguments): ) +def _check_no_partitions(value: str) -> int: + """Check that number of partitions is between [2-10,000] + Using function to validate rather than choices as error message prints all choices.""" + if value.isdigit() and 2 <= int(value) <= 10000: + return int(value) + else: + raise argparse.ArgumentTypeError( + f"{value} is not valid for number of partitions, use number in range 2 to 10000" + ) + + def _add_common_partition_arguments(optional_arguments, required_arguments): """Add all arguments common to get-partition command""" @@ -978,10 +989,9 @@ def _add_common_partition_arguments(optional_arguments, required_arguments): "--partition-num", "-pn", required=True, - help="Number of partitions/config files to generate", - type=int, - choices=range(1, 1001), - metavar="[1-1000]", + help="Number of partitions/config files to generate, a number from 2 to 10,000", + type=_check_no_partitions, + metavar="[2-10000]", ) # Optional arguments