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