src/pyutils/argparse_utils.py

   1 #!/usr/bin/python3
   2
   3 # © Copyright 2021-2023, Scott Gasch
   4
   5 """These are helpers for commandline argument parsing meant to work
   6 with Python's :mod:`argparse` module from the standard library (See:
   7 https://docs.python.org/3/library/argparse.html).  It contains
   8 validators for new argument types (such as free-form dates, durations,
   9 IP addresses, etc...)  and an action that creates a pair of flags: one
  10 to disable a feature and another to enable it.
  11
  12 See also :py:class:`pyutils.config.OptionalRawFormatter` which is
  13 automatically enabled if you use :py:mod:`config` module.
  14
  15 """
  16
  17 import argparse
  18 import datetime
  19 import logging
  20 import os
  21 from typing import Any, Optional
  22
  23 from overrides import overrides
  24
  25 # This module is commonly used by others in here and should avoid
  26 # taking any unnecessary dependencies back on them.
  27
  28 logger = logging.getLogger(__name__)
  29
  30
  31 class ActionNoYes(argparse.Action):
  32     """An argparse Action that allows for commandline arguments like this::
  33
  34         cfg.add_argument(
  35             '--enable_the_thing',
  36             action=ActionNoYes,
  37             default=False,
  38             help='Should we enable the thing?'
  39         )
  40
  41     This creates the following cmdline arguments::
  42
  43         --enable_the_thing
  44         --no_enable_the_thing
  45
  46     These arguments can be used to indicate the inclusion or exclusion of
  47     binary exclusive behaviors.
  48
  49     Raises:
  50         ValueError: illegal argument value or combination
  51
  52     """
  53
  54     def __init__(
  55         self,
  56         option_strings: str,
  57         dest: str,
  58         default: Optional[str] = None,
  59         required: bool = False,
  60         help: Optional[str] = None,
  61     ):
  62         if default is None:
  63             msg = "You must provide a default with Yes/No action"
  64             logger.critical(msg)
  65             raise ValueError(msg)
  66         if len(option_strings) != 1:
  67             msg = "Only single argument is allowed with NoYes action"
  68             logger.critical(msg)
  69             raise ValueError(msg)
  70         opt = option_strings[0]
  71         if not opt.startswith("--"):
  72             msg = "Yes/No arguments must be prefixed with --"
  73             logger.critical(msg)
  74             raise ValueError(msg)
  75
  76         opt = opt[2:]
  77         opts = ["--" + opt, "--no_" + opt]
  78         super().__init__(
  79             opts,
  80             dest,
  81             nargs=0,
  82             const=None,
  83             default=default,
  84             required=required,
  85             help=help,
  86         )
  87
  88     @overrides
  89     def __call__(self, parser, namespace, values, option_strings: Optional[str] = None):
  90         if option_strings is not None:
  91             if option_strings.startswith("--no-") or option_strings.startswith("--no_"):
  92                 setattr(namespace, self.dest, False)
  93             else:
  94                 setattr(namespace, self.dest, True)
  95
  96
  97 def valid_bool(v: Any) -> bool:
  98     """
  99     If the string is a valid bool, return its value.  Otherwise raise.
 100
 101     Args:
 102         v: data passed to an argument expecting a bool on the cmdline.
 103
 104     Returns:
 105         The boolean value of v
 106
 107     Raises:
 108         ArgumentTypeError: parse error (e.g. not a valid bool string)
 109
 110     Sample usage::
 111
 112         args.add_argument(
 113             '--auto',
 114             type=argparse_utils.valid_bool,
 115             default=None,
 116             metavar='True|False',
 117             help='Use your best judgement about --primary and --secondary',
 118         )
 119
 120     >>> valid_bool(True)
 121     True
 122
 123     >>> valid_bool("true")
 124     True
 125
 126     >>> valid_bool("yes")
 127     True
 128
 129     >>> valid_bool("on")
 130     True
 131
 132     >>> valid_bool("1")
 133     True
 134
 135     >>> valid_bool("off")   # Note: expect False; invalid would raise.
 136     False
 137
 138     >>> valid_bool(12345)
 139     Traceback (most recent call last):
 140     ...
 141     argparse.ArgumentTypeError: 12345
 142
 143     """
 144     if isinstance(v, bool):
 145         return v
 146     from pyutils.string_utils import to_bool
 147
 148     try:
 149         return to_bool(v)
 150     except Exception as e:
 151         raise argparse.ArgumentTypeError(v) from e
 152
 153
 154 def valid_ip(ip: str) -> str:
 155     """
 156     If the string is a valid IPv4 address, return it.  Otherwise raise
 157     an ArgumentTypeError.
 158
 159     Args:
 160         ip: data passed to a commandline arg expecting an IP(v4) address.
 161
 162     Returns:
 163         The IP address, if valid.
 164
 165     Raises:
 166         ArgumentTypeError: parse error (e.g. not a valid IP address string)
 167
 168     Sample usage::
 169
 170         args.add_argument(
 171             "-i",
 172             "--ip_address",
 173             metavar="TARGET_IP_ADDRESS",
 174             help="Target IP Address",
 175             type=argparse_utils.valid_ip,
 176         )
 177
 178     >>> valid_ip("1.2.3.4")
 179     '1.2.3.4'
 180
 181     >>> valid_ip("localhost")
 182     Traceback (most recent call last):
 183     ...
 184     argparse.ArgumentTypeError: localhost is an invalid IP address
 185
 186     """
 187     from pyutils.string_utils import extract_ip_v4
 188
 189     s = extract_ip_v4(ip.strip())
 190     if s is not None:
 191         return s
 192     msg = f"{ip} is an invalid IP address"
 193     logger.error(msg)
 194     raise argparse.ArgumentTypeError(msg)
 195
 196
 197 def valid_mac(mac: str) -> str:
 198     """
 199     If the string is a valid MAC address, return it.  Otherwise raise
 200     an ArgumentTypeError.
 201
 202     Args:
 203         mac: a value passed to a commandline flag expecting a MAC address.
 204
 205     Returns:
 206         The MAC address passed or raises ArgumentTypeError on error.
 207
 208     Raises:
 209         ArgumentTypeError: parse error (e.g. not a valid MAC address)
 210
 211     Sample usage::
 212
 213         group.add_argument(
 214             "-m",
 215             "--mac",
 216             metavar="MAC_ADDRESS",
 217             help="Target MAC Address",
 218             type=argparse_utils.valid_mac,
 219         )
 220
 221     >>> valid_mac('12:23:3A:4F:55:66')
 222     '12:23:3A:4F:55:66'
 223
 224     >>> valid_mac('12-23-3A-4F-55-66')
 225     '12-23-3A-4F-55-66'
 226
 227     >>> valid_mac('big')
 228     Traceback (most recent call last):
 229     ...
 230     argparse.ArgumentTypeError: big is an invalid MAC address
 231
 232     """
 233     from pyutils.string_utils import extract_mac_address
 234
 235     s = extract_mac_address(mac)
 236     if s is not None:
 237         return s
 238     msg = f"{mac} is an invalid MAC address"
 239     logger.error(msg)
 240     raise argparse.ArgumentTypeError(msg)
 241
 242
 243 def valid_percentage(num: str) -> float:
 244     """
 245     If the string is a valid (0 <= n <= 100) percentage, return it.
 246     Otherwise raise an ArgumentTypeError.
 247
 248     Arg:
 249         num: data passed to a flag expecting a percentage with a value
 250              between 0 and 100 inclusive.
 251
 252     Returns:
 253         The number if valid, otherwise raises ArgumentTypeError.
 254
 255     Raises:
 256         ArgumentTypeError: parse error (e.g. not a valid percentage)
 257
 258     Sample usage::
 259
 260         args.add_argument(
 261             '--percent_change',
 262             type=argparse_utils.valid_percentage,
 263             default=0,
 264             help='The percent change (0<=n<=100) of foobar',
 265         )
 266
 267     >>> valid_percentage("15%")
 268     15.0
 269
 270     >>> valid_percentage('40')
 271     40.0
 272
 273     >>> valid_percentage('115')
 274     Traceback (most recent call last):
 275     ...
 276     argparse.ArgumentTypeError: 115 is an invalid percentage; expected 0 <= n <= 100.0
 277
 278     """
 279     num = num.strip("%")
 280     n = float(num)
 281     if 0.0 <= n <= 100.0:
 282         return n
 283     msg = f"{num} is an invalid percentage; expected 0 <= n <= 100.0"
 284     logger.error(msg)
 285     raise argparse.ArgumentTypeError(msg)
 286
 287
 288 def valid_filename(filename: str) -> str:
 289     """
 290     If the string contains a valid filename that exists on the filesystem,
 291     return it.  Otherwise raise an ArgumentTypeError.
 292
 293     .. note::
 294
 295         This method will accept directories that exist on the filesystem
 296         in addition to files.
 297
 298     Args:
 299         filename: data passed to a flag expecting a valid filename.
 300
 301     Returns:
 302         The filename if valid, otherwise raises ArgumentTypeError.
 303
 304     Raises:
 305         ArgumentTypeError: parse error (e.g. file doesn't exist)
 306
 307     Sample usage::
 308
 309         args.add_argument(
 310             '--network_mac_addresses_file',
 311             default='/home/scott/bin/network_mac_addresses.txt',
 312             metavar='FILENAME',
 313             help='Location of the network_mac_addresses file (must exist!).',
 314             type=argparse_utils.valid_filename,
 315         )
 316
 317     >>> valid_filename('/tmp')
 318     '/tmp'
 319
 320     >>> valid_filename('wfwefwefwefwefwefwefwefwef')
 321     Traceback (most recent call last):
 322     ...
 323     argparse.ArgumentTypeError: wfwefwefwefwefwefwefwefwef was not found and is therefore invalid.
 324
 325     """
 326     s = filename.strip()
 327     if os.path.exists(s):
 328         return s
 329     msg = f"{filename} was not found and is therefore invalid."
 330     logger.error(msg)
 331     raise argparse.ArgumentTypeError(msg)
 332
 333
 334 def valid_date(txt: str) -> datetime.date:
 335     """If the string is a valid date, return it.  Otherwise raise
 336     an ArgumentTypeError.
 337
 338     Args:
 339         txt: data passed to a commandline flag expecting a date.
 340
 341     Returns:
 342         the datetime.date described by txt or raises ArgumentTypeError on error.
 343
 344     Raises:
 345         ArgumentTypeError: parse error (e.g. date not valid)
 346
 347     Sample usage::
 348
 349         cfg.add_argument(
 350             "--date",
 351             nargs=1,
 352             type=argparse_utils.valid_date,
 353             metavar="DATE STRING",
 354             default=None
 355         )
 356
 357     >>> valid_date('6/5/2021')
 358     datetime.date(2021, 6, 5)
 359
 360     .. note::
 361         dates like 'next wednesday' work fine, they are just
 362         hard to doctest for without knowing when the testcase will be
 363         executed...  See :py:mod:`pyutils.datetimes.dateparse_utils`
 364         for other examples of usable expressions.
 365
 366     >>> valid_date('next wednesday') # doctest: +ELLIPSIS
 367     -ANYTHING-
 368     """
 369     from pyutils.string_utils import to_date
 370
 371     date = to_date(txt)
 372     if date is not None:
 373         return date
 374     msg = f"Cannot parse argument as a date: {txt}"
 375     logger.error(msg)
 376     raise argparse.ArgumentTypeError(msg)
 377
 378
 379 def valid_datetime(txt: str) -> datetime.datetime:
 380     """If the string is a valid datetime, return it.  Otherwise raise
 381     an ArgumentTypeError.
 382
 383     Args:
 384         txt: data passed to a commandline flag expecting a valid datetime.datetime.
 385
 386     Returns:
 387         The datetime.datetime described by txt or raises ArgumentTypeError on error.
 388
 389     Raises:
 390         ArgumentTypeError: parse error (e.g. invalid datetime string)
 391
 392     Sample usage::
 393
 394         cfg.add_argument(
 395             "--override_timestamp",
 396             nargs=1,
 397             type=argparse_utils.valid_datetime,
 398             help="Don't use the current datetime, override to argument.",
 399             metavar="DATE/TIME STRING",
 400             default=None,
 401         )
 402
 403     >>> valid_datetime('6/5/2021 3:01:02')
 404     datetime.datetime(2021, 6, 5, 3, 1, 2)
 405
 406     >>> valid_datetime('Sun Dec 11 11:50:00 PST 2022')
 407     datetime.datetime(2022, 12, 11, 11, 50)
 408
 409     .. note::
 410         Because this code uses an English date-expression parsing grammar
 411         internally, much more complex datetimes can be expressed in free form.
 412         See :mod:`pyutils.datetimes.dateparse_utils` for details.  These
 413         are not included in here because they are hard to write valid doctests
 414         for!
 415
 416     >>> valid_datetime('next christmas at 4:15am') # doctest: +ELLIPSIS
 417     -ANYTHING-
 418     """
 419     from pyutils.string_utils import to_datetime
 420
 421     dt = to_datetime(txt)
 422     if dt is not None:
 423         return dt
 424
 425     # Don't choke on the default format of unix date.
 426     try:
 427         return datetime.datetime.strptime(txt, "%a %b %d %H:%M:%S %Z %Y")
 428     except Exception:
 429         pass
 430
 431     msg = f"Cannot parse argument as datetime: {txt}"
 432     logger.error(msg)
 433     raise argparse.ArgumentTypeError(msg)
 434
 435
 436 def valid_duration(txt: str) -> datetime.timedelta:
 437     """If the string is a valid time duration, return a
 438     datetime.timedelta representing the duration described.
 439     This uses `datetime_utils.parse_duration` to parse durations
 440     and expects data such as:
 441
 442         - 15 days, 3 hours, 15 minutes
 443         - 15 days 3 hours 15 minutes
 444         - 15d 3h 15m
 445         - 15d3h5m
 446         - 3m 2s
 447         - 1000s
 448
 449     If the duration is not parsable, raise an ArgumentTypeError.
 450
 451     Args:
 452         txt: data passed to a commandline arg expecting a duration.
 453
 454     Returns:
 455         The datetime.timedelta described by txt or raises ArgumentTypeError
 456         on error.
 457
 458     Raises:
 459         ArgumentTypeError: parse error (e.g. invalid duration string)
 460
 461     Sample usage::
 462
 463         cfg.add_argument(
 464             '--ip_cache_max_staleness',
 465             type=argparse_utils.valid_duration,
 466             default=datetime.timedelta(seconds=60 * 60 * 4),
 467             metavar='DURATION',
 468             help='Max acceptable age of the IP address cache'
 469         )
 470
 471     >>> valid_duration('15d3h5m')
 472     datetime.timedelta(days=15, seconds=11100)
 473
 474     >>> valid_duration('15 days 3 hours 5 min')
 475     datetime.timedelta(days=15, seconds=11100)
 476
 477     >>> valid_duration('3m')
 478     datetime.timedelta(seconds=180)
 479
 480     >>> valid_duration('3 days, 2 hours')
 481     datetime.timedelta(days=3, seconds=7200)
 482
 483     >>> valid_duration('a little while')
 484     Traceback (most recent call last):
 485     ...
 486     argparse.ArgumentTypeError: a little while is not a valid duration.
 487     """
 488     from pyutils.datetimes.datetime_utils import parse_duration
 489
 490     try:
 491         secs = parse_duration(txt, raise_on_error=True)
 492         return datetime.timedelta(seconds=secs)
 493     except Exception as e:
 494         logger.exception("Exception while parsing a supposed duration: %s", txt)
 495         raise argparse.ArgumentTypeError(e) from e
 496
 497
 498 def valid_byte_count(txt: str) -> int:
 499     """If the string is a valid number of bytes, return an integer
 500     representing the requested byte count.  This method uses
 501     :meth:`string_utils.suffix_string_to_number` to parse and and
 502     accepts / understands:
 503
 504         - plain numbers (123456)
 505         - numbers with ISO suffixes (Mb, Gb, Pb, etc...)
 506
 507     Args:
 508         txt: data passed to a commandline arg expecting a duration.
 509
 510     Returns:
 511         An integer number of bytes or raises ArgumentTypeError on
 512         error.
 513
 514     Raises:
 515         ArgumentTypeError: parse error (e.g. byte count not parsable)
 516
 517     Sample usage::
 518
 519         cfg.add_argument(
 520             '--max_file_size',
 521             type=argparse_utils.valid_byte_count,
 522             default=(1024 * 1024),
 523             metavar='NUM_BYTES',
 524             help='The largest file we may create',
 525         )
 526
 527     >>> valid_byte_count('1Mb')
 528     1048576
 529
 530     >>> valid_byte_count("1234567")
 531     1234567
 532
 533     >>> valid_byte_count("1M")
 534     1048576
 535
 536     >>> valid_byte_count("1.2Gb")
 537     1288490188
 538
 539     >>> valid_byte_count('1.2')      # <--- contains a decimal
 540     Traceback (most recent call last):
 541     ...
 542     argparse.ArgumentTypeError: Invalid byte count: 1.2
 543
 544     >>> valid_byte_count(1234567)    # <--- not a string
 545     Traceback (most recent call last):
 546     ...
 547     argparse.ArgumentTypeError: Invalid byte count: 1234567
 548
 549     >>> valid_byte_count('On a dark and stormy night')
 550     Traceback (most recent call last):
 551     ...
 552     argparse.ArgumentTypeError: Invalid byte count: On a dark and stormy night
 553
 554     """
 555     from pyutils.string_utils import suffix_string_to_number
 556
 557     try:
 558         num_bytes = suffix_string_to_number(txt)
 559         if num_bytes:
 560             return num_bytes
 561         raise argparse.ArgumentTypeError(f"Invalid byte count: {txt}")
 562     except Exception as e:
 563         logger.exception("Exception while parsing a supposed byte count: %s", txt)
 564         raise argparse.ArgumentTypeError(e) from e
 565
 566
 567 if __name__ == "__main__":
 568     import doctest
 569
 570     doctest.ELLIPSIS_MARKER = "-ANYTHING-"
 571     doctest.testmod()