argparse_utils.py

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