third_party/pigweed/repo/pw_cli/py/pw_cli/envparse.py

   1 # Copyright 2020 The Pigweed Authors
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
   4 # use this file except in compliance with the License. You may obtain a copy of
   5 # the License at
   6 #
   7 #     https://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  11 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  12 # License for the specific language governing permissions and limitations under
  13 # the License.
  14 """The envparse module defines an environment variable parser."""
  15
  16 import argparse
  17 import os
  18 from typing import Callable, Dict, Generic, IO, List, Mapping
  19 from typing import NamedTuple, Optional, TypeVar
  20
  21
  22 class EnvNamespace(argparse.Namespace):  # pylint: disable=too-few-public-methods
  23     """Base class for parsed environment variable namespaces."""
  24
  25
  26 T = TypeVar('T')
  27 TypeConversion = Callable[[str], T]
  28
  29
  30 class VariableDescriptor(NamedTuple, Generic[T]):
  31     name: str
  32     type: TypeConversion[T]
  33     default: Optional[T]
  34
  35
  36 class EnvironmentValueError(Exception):
  37     """Exception indicating a bad type conversion on an environment variable.
  38
  39     Stores a reference to the lower-level exception from the type conversion
  40     function through the __cause__ attribute for more detailed information on
  41     the error.
  42     """
  43     def __init__(self, variable: str, value: str):
  44         self.variable: str = variable
  45         self.value: str = value
  46         super().__init__(
  47             f'Bad value for environment variable {variable}: {value}')
  48
  49
  50 class EnvironmentParser:
  51     """Parser for environment variables.
  52
  53     Args:
  54         prefix: If provided, checks that all registered environment variables
  55           start with the specified string.
  56         error_on_unrecognized: If True and prefix is provided, will raise an
  57           exception if the environment contains a variable with the specified
  58           prefix that is not registered on the EnvironmentParser.
  59
  60     Example:
  61
  62         parser = envparse.EnvironmentParser(prefix='PW_')
  63         parser.add_var('PW_LOG_LEVEL')
  64         parser.add_var('PW_LOG_FILE', type=envparse.FileType('w'))
  65         parser.add_var('PW_USE_COLOR', type=envparse.strict_bool, default=False)
  66         env = parser.parse_env()
  67
  68         configure_logging(env.PW_LOG_LEVEL, env.PW_LOG_FILE)
  69     """
  70     def __init__(self,
  71                  prefix: Optional[str] = None,
  72                  error_on_unrecognized: bool = True) -> None:
  73         self._prefix: Optional[str] = prefix
  74         self._error_on_unrecognized: bool = error_on_unrecognized
  75         self._variables: Dict[str, VariableDescriptor] = {}
  76         self._allowed_suffixes: List[str] = []
  77
  78     def add_var(
  79         self,
  80         name: str,
  81         # pylint: disable=redefined-builtin
  82         type: TypeConversion[T] = str,  # type: ignore
  83         # pylint: enable=redefined-builtin
  84         default: Optional[T] = None,
  85     ) -> None:
  86         """Registers an environment variable.
  87
  88         Args:
  89             name: The environment variable's name.
  90             type: Type conversion for the variable's value.
  91             default: Default value for the variable.
  92
  93         Raises:
  94             ValueError: If prefix was provided to the constructor and name does
  95               not start with the prefix.
  96         """
  97         if self._prefix is not None and not name.startswith(self._prefix):
  98             raise ValueError(
  99                 f'Variable {name} does not have prefix {self._prefix}')
 100
 101         self._variables[name] = VariableDescriptor(
 102             name,
 103             type,  # type: ignore
 104             default)  # type: ignore
 105
 106     def add_allowed_suffix(self, suffix: str) -> None:
 107         """Registers an environmant variable name suffix to be allowed."""
 108
 109         self._allowed_suffixes.append(suffix)
 110
 111     def parse_env(self,
 112                   env: Optional[Mapping[str, str]] = None) -> EnvNamespace:
 113         """Parses known environment variables into a namespace.
 114
 115         Args:
 116             env: Dictionary of environment variables. Defaults to os.environ.
 117
 118         Raises:
 119             EnvironmentValueError: If the type conversion fails.
 120         """
 121         if env is None:
 122             env = os.environ
 123
 124         namespace = EnvNamespace()
 125
 126         for var, desc in self._variables.items():
 127             if var not in env:
 128                 val = desc.default
 129             else:
 130                 try:
 131                     val = desc.type(env[var])
 132                 except Exception as err:
 133                     raise EnvironmentValueError(var, env[var]) from err
 134
 135             setattr(namespace, var, val)
 136
 137         allowed_suffixes = tuple(self._allowed_suffixes)
 138         for var in env:
 139             if (not hasattr(namespace, var)
 140                     and (self._prefix is None or var.startswith(self._prefix))
 141                     and var.endswith(allowed_suffixes)):
 142                 setattr(namespace, var, env[var])
 143
 144         if self._prefix is not None and self._error_on_unrecognized:
 145             for var in env:
 146                 if (var.startswith(self._prefix) and var not in self._variables
 147                         and not var.endswith(allowed_suffixes)):
 148                     raise ValueError(
 149                         f'Unrecognized environment variable {var}')
 150
 151         return namespace
 152
 153     def __repr__(self) -> str:
 154         return f'{type(self).__name__}(prefix={self._prefix})'
 155
 156
 157 # List of emoji which are considered to represent "True".
 158 _BOOLEAN_TRUE_EMOJI = set([
 159     '✔️',
 160     '👍',
 161     '👍🏻',
 162     '👍🏼',
 163     '👍🏽',
 164     '👍🏾',
 165     '👍🏿',
 166     '💯',
 167 ])
 168
 169
 170 def strict_bool(value: str) -> bool:
 171     return (value == '1' or value.lower() == 'true'
 172             or value in _BOOLEAN_TRUE_EMOJI)
 173
 174
 175 # TODO(mohrr) Switch to Literal when no longer supporting Python 3.7.
 176 # OpenMode = Literal['r', 'rb', 'w', 'wb']
 177 OpenMode = str
 178
 179
 180 class FileType:
 181     def __init__(self, mode: OpenMode) -> None:
 182         self._mode: OpenMode = mode
 183
 184     def __call__(self, value: str) -> IO:
 185         return open(value, self._mode)