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 16import argparse 17from dataclasses import dataclass 18import os 19from typing import ( 20 Callable, 21 Dict, 22 Generic, 23 IO, 24 List, 25 Mapping, 26 Optional, 27 TypeVar, 28 Union, 29) 30 31 32class EnvNamespace( 33 argparse.Namespace 34): # pylint: disable=too-few-public-methods 35 """Base class for parsed environment variable namespaces.""" 36 37 38T = TypeVar('T') 39TypeConversion = Callable[[str], T] 40 41 42@dataclass 43class VariableDescriptor(Generic[T]): 44 name: str 45 type: TypeConversion[T] 46 default: Optional[T] 47 48 49class EnvironmentValueError(Exception): 50 """Exception indicating a bad type conversion on an environment variable. 51 52 Stores a reference to the lower-level exception from the type conversion 53 function through the __cause__ attribute for more detailed information on 54 the error. 55 """ 56 57 def __init__(self, variable: str, value: str): 58 self.variable: str = variable 59 self.value: str = value 60 super().__init__( 61 f'Bad value for environment variable {variable}: {value}' 62 ) 63 64 65class EnvironmentParser: 66 """Parser for environment variables. 67 68 Args: 69 prefix: If provided, checks that all registered environment variables 70 start with the specified string. 71 error_on_unrecognized: If True and prefix is provided, will raise an 72 exception if the environment contains a variable with the specified 73 prefix that is not registered on the EnvironmentParser. If None, 74 checks existence of PW_ENVIRONMENT_NO_ERROR_ON_UNRECOGNIZED (but not 75 value). 76 77 Example: 78 79 parser = envparse.EnvironmentParser(prefix='PW_') 80 parser.add_var('PW_LOG_LEVEL') 81 parser.add_var('PW_LOG_FILE', type=envparse.FileType('w')) 82 parser.add_var('PW_USE_COLOR', type=envparse.strict_bool, default=False) 83 env = parser.parse_env() 84 85 configure_logging(env.PW_LOG_LEVEL, env.PW_LOG_FILE) 86 """ 87 88 def __init__( 89 self, 90 prefix: Optional[str] = None, 91 error_on_unrecognized: Union[bool, None] = None, 92 ) -> None: 93 self._prefix: Optional[str] = prefix 94 if error_on_unrecognized is None: 95 varname = 'PW_ENVIRONMENT_NO_ERROR_ON_UNRECOGNIZED' 96 error_on_unrecognized = varname not in os.environ 97 self._error_on_unrecognized: bool = error_on_unrecognized 98 99 self._variables: Dict[str, VariableDescriptor] = {} 100 self._allowed_suffixes: List[str] = [] 101 102 def add_var( 103 self, 104 name: str, 105 # pylint: disable=redefined-builtin 106 type: TypeConversion[T] = str, # type: ignore[assignment] 107 # pylint: enable=redefined-builtin 108 default: Optional[T] = None, 109 ) -> None: 110 """Registers an environment variable. 111 112 Args: 113 name: The environment variable's name. 114 type: Type conversion for the variable's value. 115 default: Default value for the variable. 116 117 Raises: 118 ValueError: If prefix was provided to the constructor and name does 119 not start with the prefix. 120 """ 121 if self._prefix is not None and not name.startswith(self._prefix): 122 raise ValueError( 123 f'Variable {name} does not have prefix {self._prefix}' 124 ) 125 126 self._variables[name] = VariableDescriptor(name, type, default) 127 128 def add_allowed_suffix(self, suffix: str) -> None: 129 """Registers an environment variable name suffix to be allowed.""" 130 131 self._allowed_suffixes.append(suffix) 132 133 def parse_env( 134 self, env: Optional[Mapping[str, str]] = None 135 ) -> EnvNamespace: 136 """Parses known environment variables into a namespace. 137 138 Args: 139 env: Dictionary of environment variables. Defaults to os.environ. 140 141 Raises: 142 EnvironmentValueError: If the type conversion fails. 143 """ 144 if env is None: 145 env = os.environ 146 147 namespace = EnvNamespace() 148 149 for var, desc in self._variables.items(): 150 if var not in env: 151 val = desc.default 152 else: 153 try: 154 val = desc.type(env[var]) # type: ignore 155 except Exception as err: 156 raise EnvironmentValueError(var, env[var]) from err 157 158 setattr(namespace, var, val) 159 160 allowed_suffixes = tuple(self._allowed_suffixes) 161 for var in env: 162 if ( 163 not hasattr(namespace, var) 164 and (self._prefix is None or var.startswith(self._prefix)) 165 and var.endswith(allowed_suffixes) 166 ): 167 setattr(namespace, var, env[var]) 168 169 if self._prefix is not None and self._error_on_unrecognized: 170 for var in env: 171 if ( 172 var.startswith(self._prefix) 173 and var not in self._variables 174 and not var.endswith(allowed_suffixes) 175 ): 176 raise ValueError(f'Unrecognized environment variable {var}') 177 178 return namespace 179 180 def __repr__(self) -> str: 181 return f'{type(self).__name__}(prefix={self._prefix})' 182 183 184# List of emoji which are considered to represent "True". 185_BOOLEAN_TRUE_EMOJI = set( 186 [ 187 '✔️', 188 '��', 189 '����', 190 '����', 191 '����', 192 '����', 193 '����', 194 '��', 195 ] 196) 197 198 199def strict_bool(value: str) -> bool: 200 return ( 201 value == '1' or value.lower() == 'true' or value in _BOOLEAN_TRUE_EMOJI 202 ) 203 204 205# TODO(mohrr) Switch to Literal when no longer supporting Python 3.7. 206# OpenMode = Literal['r', 'rb', 'w', 'wb'] 207OpenMode = str 208 209 210class FileType: 211 def __init__(self, mode: OpenMode) -> None: 212 self._mode: OpenMode = mode 213 214 def __call__(self, value: str) -> IO: 215 return open(value, self._mode) 216