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