argparse_utils.py

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