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