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