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 cmdline arguments:
  32
  33         --enable_the_thing
  34         --no_enable_the_thing
  35
  36     """
  37
  38     def __init__(self, option_strings, dest, default=None, required=False, help=None):
  39         if default is None:
  40             msg = 'You must provide a default with Yes/No action'
  41             logger.critical(msg)
  42             raise ValueError(msg)
  43         if len(option_strings) != 1:
  44             msg = 'Only single argument is allowed with NoYes action'
  45             logger.critical(msg)
  46             raise ValueError(msg)
  47         opt = option_strings[0]
  48         if not opt.startswith('--'):
  49             msg = 'Yes/No arguments must be prefixed with --'
  50             logger.critical(msg)
  51             raise ValueError(msg)
  52
  53         opt = opt[2:]
  54         opts = ['--' + opt, '--no_' + opt]
  55         super().__init__(
  56             opts,
  57             dest,
  58             nargs=0,
  59             const=None,
  60             default=default,
  61             required=required,
  62             help=help,
  63         )
  64
  65     @overrides
  66     def __call__(self, parser, namespace, values, option_strings=None):
  67         if option_strings.startswith('--no-') or option_strings.startswith('--no_'):
  68             setattr(namespace, self.dest, False)
  69         else:
  70             setattr(namespace, self.dest, True)
  71
  72
  73 def valid_bool(v: Any) -> bool:
  74     """
  75     If the string is a valid bool, return its value.
  76
  77     >>> valid_bool(True)
  78     True
  79
  80     >>> valid_bool("true")
  81     True
  82
  83     >>> valid_bool("yes")
  84     True
  85
  86     >>> valid_bool("on")
  87     True
  88
  89     >>> valid_bool("1")
  90     True
  91
  92     >>> valid_bool(12345)
  93     Traceback (most recent call last):
  94     ...
  95     argparse.ArgumentTypeError: 12345
  96
  97     """
  98     if isinstance(v, bool):
  99         return v
 100     from string_utils import to_bool
 101
 102     try:
 103         return to_bool(v)
 104     except Exception as e:
 105         raise argparse.ArgumentTypeError(v) from e
 106
 107
 108 def valid_ip(ip: str) -> str:
 109     """
 110     If the string is a valid IPv4 address, return it.  Otherwise raise
 111     an ArgumentTypeError.
 112
 113     >>> valid_ip("1.2.3.4")
 114     '1.2.3.4'
 115
 116     >>> valid_ip("localhost")
 117     Traceback (most recent call last):
 118     ...
 119     argparse.ArgumentTypeError: localhost is an invalid IP address
 120
 121     """
 122     from string_utils import extract_ip_v4
 123
 124     s = extract_ip_v4(ip.strip())
 125     if s is not None:
 126         return s
 127     msg = f"{ip} is an invalid IP address"
 128     logger.error(msg)
 129     raise argparse.ArgumentTypeError(msg)
 130
 131
 132 def valid_mac(mac: str) -> str:
 133     """
 134     If the string is a valid MAC address, return it.  Otherwise raise
 135     an ArgumentTypeError.
 136
 137     >>> valid_mac('12:23:3A:4F:55:66')
 138     '12:23:3A:4F:55:66'
 139
 140     >>> valid_mac('12-23-3A-4F-55-66')
 141     '12-23-3A-4F-55-66'
 142
 143     >>> valid_mac('big')
 144     Traceback (most recent call last):
 145     ...
 146     argparse.ArgumentTypeError: big is an invalid MAC address
 147
 148     """
 149     from string_utils import extract_mac_address
 150
 151     s = extract_mac_address(mac)
 152     if s is not None:
 153         return s
 154     msg = f"{mac} is an invalid MAC address"
 155     logger.error(msg)
 156     raise argparse.ArgumentTypeError(msg)
 157
 158
 159 def valid_percentage(num: str) -> float:
 160     """
 161     If the string is a valid percentage, return it.  Otherwise raise
 162     an ArgumentTypeError.
 163
 164     >>> valid_percentage("15%")
 165     15.0
 166
 167     >>> valid_percentage('40')
 168     40.0
 169
 170     >>> valid_percentage('115')
 171     Traceback (most recent call last):
 172     ...
 173     argparse.ArgumentTypeError: 115 is an invalid percentage; expected 0 <= n <= 100.0
 174
 175     """
 176     num = num.strip('%')
 177     n = float(num)
 178     if 0.0 <= n <= 100.0:
 179         return n
 180     msg = f"{num} is an invalid percentage; expected 0 <= n <= 100.0"
 181     logger.error(msg)
 182     raise argparse.ArgumentTypeError(msg)
 183
 184
 185 def valid_filename(filename: str) -> str:
 186     """
 187     If the string is a valid filename, return it.  Otherwise raise
 188     an ArgumentTypeError.
 189
 190     >>> valid_filename('/tmp')
 191     '/tmp'
 192
 193     >>> valid_filename('wfwefwefwefwefwefwefwefwef')
 194     Traceback (most recent call last):
 195     ...
 196     argparse.ArgumentTypeError: wfwefwefwefwefwefwefwefwef was not found and is therefore invalid.
 197
 198     """
 199     s = filename.strip()
 200     if os.path.exists(s):
 201         return s
 202     msg = f"{filename} was not found and is therefore invalid."
 203     logger.error(msg)
 204     raise argparse.ArgumentTypeError(msg)
 205
 206
 207 def valid_date(txt: str) -> datetime.date:
 208     """If the string is a valid date, return it.  Otherwise raise
 209     an ArgumentTypeError.
 210
 211     >>> valid_date('6/5/2021')
 212     datetime.date(2021, 6, 5)
 213
 214     # Note: dates like 'next wednesday' work fine, they are just
 215     # hard to test for without knowing when the testcase will be
 216     # executed...
 217     >>> valid_date('next wednesday') # doctest: +ELLIPSIS
 218     -ANYTHING-
 219     """
 220     from string_utils import to_date
 221
 222     date = to_date(txt)
 223     if date is not None:
 224         return date
 225     msg = f'Cannot parse argument as a date: {txt}'
 226     logger.error(msg)
 227     raise argparse.ArgumentTypeError(msg)
 228
 229
 230 def valid_datetime(txt: str) -> datetime.datetime:
 231     """If the string is a valid datetime, return it.  Otherwise raise
 232     an ArgumentTypeError.
 233
 234     >>> valid_datetime('6/5/2021 3:01:02')
 235     datetime.datetime(2021, 6, 5, 3, 1, 2)
 236
 237     # Again, these types of expressions work fine but are
 238     # difficult to test with doctests because the answer is
 239     # relative to the time the doctest is executed.
 240     >>> valid_datetime('next christmas at 4:15am') # doctest: +ELLIPSIS
 241     -ANYTHING-
 242     """
 243     from string_utils import to_datetime
 244
 245     dt = to_datetime(txt)
 246     if dt is not None:
 247         return dt
 248     msg = f'Cannot parse argument as datetime: {txt}'
 249     logger.error(msg)
 250     raise argparse.ArgumentTypeError(msg)
 251
 252
 253 def valid_duration(txt: str) -> datetime.timedelta:
 254     """If the string is a valid time duration, return a
 255     datetime.timedelta representing the period of time.  Otherwise
 256     maybe raise an ArgumentTypeError or potentially just treat the
 257     time window as zero in length.
 258
 259     >>> valid_duration('3m')
 260     datetime.timedelta(seconds=180)
 261
 262     >>> valid_duration('your mom')
 263     datetime.timedelta(0)
 264
 265     """
 266     from datetime_utils import parse_duration
 267
 268     try:
 269         secs = parse_duration(txt)
 270         return datetime.timedelta(seconds=secs)
 271     except Exception as e:
 272         logger.exception(e)
 273         raise argparse.ArgumentTypeError(e) from e
 274
 275
 276 if __name__ == '__main__':
 277     import doctest
 278
 279     doctest.ELLIPSIS_MARKER = '-ANYTHING-'
 280     doctest.testmod()