pw_cli: add allowed suffixes to env parser
Add ability to ignore suffixes in environment variable parser. Add
'_CIPD_INSTALL_DIR' to this list.
Change-Id: I35c734136591fb932ccb7d99c58ecf3582d4e7f2
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/12842
Reviewed-by: Alexei Frolov <frolv@google.com>
Commit-Queue: Rob Mohr <mohrr@google.com>
diff --git a/pw_cli/py/pw_cli/env.py b/pw_cli/py/pw_cli/env.py
index 41f16f3..dc9eaa7 100644
--- a/pw_cli/py/pw_cli/env.py
+++ b/pw_cli/py/pw_cli/env.py
@@ -40,9 +40,7 @@
parser.add_var('PW_USE_COLOR', type=envparse.strict_bool, default=False)
parser.add_var('PW_USE_GCS_ENVSETUP', type=envparse.strict_bool)
- parser.add_var('PW_PIGWEED_CIPD_INSTALL_DIR')
- parser.add_var('PW_LUCI_CIPD_INSTALL_DIR')
- parser.add_var('PW_CIPD_INSTALL_DIR')
+ parser.add_allowed_suffix('_CIPD_INSTALL_DIR')
parser.add_var('PW_CIPD_PACKAGE_FILES')
parser.add_var('PW_VIRTUALENV_REQUIREMENTS')
diff --git a/pw_cli/py/pw_cli/envparse.py b/pw_cli/py/pw_cli/envparse.py
index c65330b..2450992 100644
--- a/pw_cli/py/pw_cli/envparse.py
+++ b/pw_cli/py/pw_cli/envparse.py
@@ -15,8 +15,8 @@
import argparse
import os
-from typing import Callable, Dict, Generic, IO, Literal, Mapping, NamedTuple
-from typing import Optional, TypeVar
+from typing import Callable, Dict, Generic, IO, List, Literal, Mapping
+from typing import NamedTuple, Optional, TypeVar
class EnvNamespace(argparse.Namespace): # pylint: disable=too-few-public-methods
@@ -73,6 +73,7 @@
self._prefix: Optional[str] = prefix
self._error_on_unrecognized: bool = error_on_unrecognized
self._variables: Dict[str, VariableDescriptor] = {}
+ self._allowed_suffixes: List[str] = []
def add_var(
self,
@@ -102,6 +103,11 @@
type, # type: ignore
default) # type: ignore
+ def add_allowed_suffix(self, suffix: str) -> None:
+ """Registers an environmant variable name suffix to be allowed."""
+
+ self._allowed_suffixes.append(suffix)
+
def parse_env(self,
env: Optional[Mapping[str, str]] = None) -> EnvNamespace:
"""Parses known environment variables into a namespace.
@@ -128,9 +134,17 @@
setattr(namespace, var, val)
+ allowed_suffixes = tuple(self._allowed_suffixes)
+ for var in env:
+ if (not hasattr(namespace, var)
+ and (self._prefix is None or var.startswith(self._prefix))
+ and var.endswith(allowed_suffixes)):
+ setattr(namespace, var, env[var])
+
if self._prefix is not None and self._error_on_unrecognized:
for var in env:
- if var.startswith(self._prefix) and var not in self._variables:
+ if (var.startswith(self._prefix) and var not in self._variables
+ and not var.endswith(allowed_suffixes)):
raise ValueError(
f'Unrecognized environment variable {var}')
diff --git a/pw_cli/py/pw_cli/envparse_test.py b/pw_cli/py/pw_cli/envparse_test.py
index 8917a42..0a6f180 100644
--- a/pw_cli/py/pw_cli/envparse_test.py
+++ b/pw_cli/py/pw_cli/envparse_test.py
@@ -96,6 +96,7 @@
'PW_FOO': '001',
'PW_BAR': '010',
'PW_BAZ': '100',
+ 'IGNORED': '011',
}
def test_parse_unrecognized_variable(self):
@@ -106,6 +107,20 @@
with self.assertRaises(ValueError):
parser.parse_env(env=self.raw_env)
+ def test_parse_unrecognized_but_allowed_suffix(self):
+ parser = envparse.EnvironmentParser(prefix='PW_')
+ parser.add_allowed_suffix('_ALLOWED_SUFFIX')
+
+ env = parser.parse_env(env={'PW_FOO_ALLOWED_SUFFIX': '001'})
+ self.assertEqual(env.PW_FOO_ALLOWED_SUFFIX, '001')
+
+ def test_parse_allowed_suffix_but_not_suffix(self):
+ parser = envparse.EnvironmentParser(prefix='PW_')
+ parser.add_allowed_suffix('_ALLOWED_SUFFIX')
+
+ with self.assertRaises(ValueError):
+ parser.parse_env(env={'PW_FOO_ALLOWED_SUFFIX_FOO': '001'})
+
def test_parse_ignore_unrecognized(self):
parser = envparse.EnvironmentParser(prefix='PW_',
error_on_unrecognized=False)
