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