src/pyutils/argparse_utils.py

   1 #!/usr/bin/python3
   2
   3 # © Copyright 2021-2022, Scott Gasch
   4
   5 """Helpers for commandline argument parsing."""
   6
   7 import argparse
   8 import datetime
   9 import logging
  10 import os
  11 from typing import Any
  12
  13 from overrides import overrides
  14
  15 # This module is commonly used by others in here and should avoid
  16 # taking any unnecessary dependencies back on them.
  17
  18 logger = logging.getLogger(__name__)
  19
  20
  21 class ActionNoYes(argparse.Action):
  22     """An argparse Action that allows for commandline arguments like this::
  23
  24         cfg.add_argument(
  25             '--enable_the_thing',
  26             action=ActionNoYes,
  27             default=False,
  28             help='Should we enable the thing?'
  29         )
  30
  31     This creates the following cmdline arguments::
  32
  33         --enable_the_thing
  34         --no_enable_the_thing
  35
  36     These arguments can be used to indicate the inclusion or exclusion of
  37     binary exclusive behaviors.
  38     """
  39
  40     def __init__(self, option_strings, dest, default=None, required=False, help=None):
  41         if default is None:
  42             msg = 'You must provide a default with Yes/No action'
  43             logger.critical(msg)
  44             raise ValueError(msg)
  45         if len(option_strings) != 1:
  46             msg = 'Only single argument is allowed with NoYes action'
  47             logger.critical(msg)
  48             raise ValueError(msg)
  49         opt = option_strings[0]
  50         if not opt.startswith('--'):
  51             msg = 'Yes/No arguments must be prefixed with --'
  52             logger.critical(msg)
  53             raise ValueError(msg)
  54
  55         opt = opt[2:]
  56         opts = ['--' + opt, '--no_' + opt]
  57         super().__init__(
  58             opts,
  59             dest,
  60             nargs=0,
  61             const=None,
  62             default=default,
  63             required=required,
  64             help=help,
  65         )
  66
  67     @overrides
  68     def __call__(self, parser, namespace, values, option_strings=None):
  69         if option_strings.startswith('--no-') or option_strings.startswith('--no_'):
  70             setattr(namespace, self.dest, False)
  71         else:
  72             setattr(namespace, self.dest, True)
  73
  74
  75 def valid_bool(v: Any) -> bool:
  76     """
  77     If the string is a valid bool, return its value.  Sample usage::
  78
  79         args.add_argument(
  80             '--auto',
  81             type=argparse_utils.valid_bool,
  82             default=None,
  83             metavar='True|False',
  84             help='Use your best judgement about --primary and --secondary',
  85         )
  86
  87     >>> valid_bool(True)
  88     True
  89
  90     >>> valid_bool("true")
  91     True
  92
  93     >>> valid_bool("yes")
  94     True
  95
  96     >>> valid_bool("on")
  97     True
  98
  99     >>> valid_bool("1")
 100     True
 101
 102     >>> valid_bool(12345)
 103     Traceback (most recent call last):
 104     ...
 105     argparse.ArgumentTypeError: 12345
 106
 107     """
 108     if isinstance(v, bool):
 109         return v
 110     from pyutils.string_utils import to_bool
 111
 112     try:
 113         return to_bool(v)
 114     except Exception as e:
 115         raise argparse.ArgumentTypeError(v) from e
 116
 117
 118 def valid_ip(ip: str) -> str:
 119     """
 120     If the string is a valid IPv4 address, return it.  Otherwise raise
 121     an ArgumentTypeError.  Sample usage::
 122
 123         group.add_argument(
 124             "-i",
 125             "--ip_address",
 126             metavar="TARGET_IP_ADDRESS",
 127             help="Target IP Address",
 128             type=argparse_utils.valid_ip,
 129         )
 130
 131     >>> valid_ip("1.2.3.4")
 132     '1.2.3.4'
 133
 134     >>> valid_ip("localhost")
 135     Traceback (most recent call last):
 136     ...
 137     argparse.ArgumentTypeError: localhost is an invalid IP address
 138
 139     """
 140     from pyutils.string_utils import extract_ip_v4
 141
 142     s = extract_ip_v4(ip.strip())
 143     if s is not None:
 144         return s
 145     msg = f"{ip} is an invalid IP address"
 146     logger.error(msg)
 147     raise argparse.ArgumentTypeError(msg)
 148
 149
 150 def valid_mac(mac: str) -> str:
 151     """
 152     If the string is a valid MAC address, return it.  Otherwise raise
 153     an ArgumentTypeError.  Sample usage::
 154
 155         group.add_argument(
 156             "-m",
 157             "--mac",
 158             metavar="MAC_ADDRESS",
 159             help="Target MAC Address",
 160             type=argparse_utils.valid_mac,
 161         )
 162
 163     >>> valid_mac('12:23:3A:4F:55:66')
 164     '12:23:3A:4F:55:66'
 165
 166     >>> valid_mac('12-23-3A-4F-55-66')
 167     '12-23-3A-4F-55-66'
 168
 169     >>> valid_mac('big')
 170     Traceback (most recent call last):
 171     ...
 172     argparse.ArgumentTypeError: big is an invalid MAC address
 173
 174     """
 175     from pyutils.string_utils import extract_mac_address
 176
 177     s = extract_mac_address(mac)
 178     if s is not None:
 179         return s
 180     msg = f"{mac} is an invalid MAC address"
 181     logger.error(msg)
 182     raise argparse.ArgumentTypeError(msg)
 183
 184
 185 def valid_percentage(num: str) -> float:
 186     """
 187     If the string is a valid (0 <= n <= 100) percentage, return it.
 188     Otherwise raise an ArgumentTypeError.  Sample usage::
 189
 190         args.add_argument(
 191             '--percent_change',
 192             type=argparse_utils.valid_percentage,
 193             default=0,
 194             help='The percent change (0<=n<=100) of foobar',
 195         )
 196
 197     >>> valid_percentage("15%")
 198     15.0
 199
 200     >>> valid_percentage('40')
 201     40.0
 202
 203     >>> valid_percentage('115')
 204     Traceback (most recent call last):
 205     ...
 206     argparse.ArgumentTypeError: 115 is an invalid percentage; expected 0 <= n <= 100.0
 207
 208     """
 209     num = num.strip('%')
 210     n = float(num)
 211     if 0.0 <= n <= 100.0:
 212         return n
 213     msg = f"{num} is an invalid percentage; expected 0 <= n <= 100.0"
 214     logger.error(msg)
 215     raise argparse.ArgumentTypeError(msg)
 216
 217
 218 def valid_filename(filename: str) -> str:
 219     """
 220     If the string is a valid filename, return it.  Otherwise raise
 221     an ArgumentTypeError.  Sample usage::
 222
 223         args.add_argument(
 224             '--network_mac_addresses_file',
 225             default='/home/scott/bin/network_mac_addresses.txt',
 226             metavar='FILENAME',
 227             help='Location of the network_mac_addresses file (must exist!).',
 228             type=argparse_utils.valid_filename,
 229         )
 230
 231     >>> valid_filename('/tmp')
 232     '/tmp'
 233
 234     >>> valid_filename('wfwefwefwefwefwefwefwefwef')
 235     Traceback (most recent call last):
 236     ...
 237     argparse.ArgumentTypeError: wfwefwefwefwefwefwefwefwef was not found and is therefore invalid.
 238
 239     """
 240     s = filename.strip()
 241     if os.path.exists(s):
 242         return s
 243     msg = f"{filename} was not found and is therefore invalid."
 244     logger.error(msg)
 245     raise argparse.ArgumentTypeError(msg)
 246
 247
 248 def valid_date(txt: str) -> datetime.date:
 249     """If the string is a valid date, return it.  Otherwise raise
 250     an ArgumentTypeError.  Sample usage::
 251
 252         cfg.add_argument(
 253             "--date",
 254             nargs=1,
 255             type=argparse_utils.valid_date,
 256             metavar="DATE STRING",
 257             default=None
 258         )
 259
 260     >>> valid_date('6/5/2021')
 261     datetime.date(2021, 6, 5)
 262
 263     Note: dates like 'next wednesday' work fine, they are just
 264     hard to test for without knowing when the testcase will be
 265     executed...
 266
 267     >>> valid_date('next wednesday') # doctest: +ELLIPSIS
 268     -ANYTHING-
 269     """
 270     from pyutils.string_utils import to_date
 271
 272     date = to_date(txt)
 273     if date is not None:
 274         return date
 275     msg = f'Cannot parse argument as a date: {txt}'
 276     logger.error(msg)
 277     raise argparse.ArgumentTypeError(msg)
 278
 279
 280 def valid_datetime(txt: str) -> datetime.datetime:
 281     """If the string is a valid datetime, return it.  Otherwise raise
 282     an ArgumentTypeError.  Sample usage::
 283
 284         cfg.add_argument(
 285             "--override_timestamp",
 286             nargs=1,
 287             type=argparse_utils.valid_datetime,
 288             help="Don't use the current datetime, override to argument.",
 289             metavar="DATE/TIME STRING",
 290             default=None,
 291         )
 292
 293     >>> valid_datetime('6/5/2021 3:01:02')
 294     datetime.datetime(2021, 6, 5, 3, 1, 2)
 295
 296     Because this thing uses an English date-expression parsing grammar
 297     internally, much more complex datetimes can be expressed in free form.
 298     See: `tests/datetimez/dateparse_utils_test.py` for examples.  These
 299     are not included in here because they are hard to write valid doctests
 300     for!
 301
 302     >>> valid_datetime('next christmas at 4:15am') # doctest: +ELLIPSIS
 303     -ANYTHING-
 304     """
 305     from pyutils.string_utils import to_datetime
 306
 307     dt = to_datetime(txt)
 308     if dt is not None:
 309         return dt
 310     msg = f'Cannot parse argument as datetime: {txt}'
 311     logger.error(msg)
 312     raise argparse.ArgumentTypeError(msg)
 313
 314
 315 def valid_duration(txt: str) -> datetime.timedelta:
 316     """If the string is a valid time duration, return a
 317     datetime.timedelta representing the period of time.
 318     Sample usage::
 319
 320         cfg.add_argument(
 321             '--ip_cache_max_staleness',
 322             type=argparse_utils.valid_duration,
 323             default=datetime.timedelta(seconds=60 * 60 * 12),
 324             metavar='DURATION',
 325             help='Max acceptable age of the IP address cache'
 326         )
 327
 328     >>> valid_duration('3m')
 329     datetime.timedelta(seconds=180)
 330
 331     >>> valid_duration('3 days, 2 hours')
 332     datetime.timedelta(days=3, seconds=7200)
 333
 334     >>> valid_duration('a little while')
 335     Traceback (most recent call last):
 336     ...
 337     argparse.ArgumentTypeError: a little while is not a valid duration.
 338
 339     """
 340     from pyutils.datetimez.datetime_utils import parse_duration
 341
 342     try:
 343         secs = parse_duration(txt, raise_on_error=True)
 344         return datetime.timedelta(seconds=secs)
 345     except Exception as e:
 346         logger.exception(e)
 347         raise argparse.ArgumentTypeError(e) from e
 348
 349
 350 if __name__ == '__main__':
 351     import doctest
 352
 353     doctest.ELLIPSIS_MARKER = '-ANYTHING-'
 354     doctest.testmod()