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