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