diff options
author | Ross Burton <ross@burtonini.com> | 2021-08-10 17:55:06 +0100 |
---|---|---|
committer | Richard Purdie <richard.purdie@linuxfoundation.org> | 2021-08-12 06:27:39 +0100 |
commit | 9974848f67581ff7d76cef52a94f505af99b4932 (patch) | |
tree | ff4c03dc89b98843e20a783cfa9fb6596e50fc49 /lib | |
parent | 06a0bbe746f879ae539223e7fdb6f07d55d13719 (diff) | |
download | bitbake-9974848f67581ff7d76cef52a94f505af99b4932.tar.gz |
utils: add environment updating context manager
bb.utils.environment() is a context manager to alter os.environ inside
a specific block, restoring it after the block is closed.
Signed-off-by: Ross Burton <ross.burton@arm.com>
Signed-off-by: Richard Purdie <richard.purdie@linuxfoundation.org>
Diffstat (limited to 'lib')
-rw-r--r-- | lib/bb/tests/utils.py | 18 | ||||
-rw-r--r-- | lib/bb/utils.py | 16 |
2 files changed, 34 insertions, 0 deletions
diff --git a/lib/bb/tests/utils.py b/lib/bb/tests/utils.py index a7ff33db5..4d5e21b99 100644 --- a/lib/bb/tests/utils.py +++ b/lib/bb/tests/utils.py @@ -666,3 +666,21 @@ class GetReferencedVars(unittest.TestCase): layers = [{"SRC_URI"}, {"QT_GIT", "QT_MODULE", "QT_MODULE_BRANCH_PARAM", "QT_GIT_PROTOCOL"}, {"QT_GIT_PROJECT", "QT_MODULE_BRANCH", "BPN"}, {"PN", "SPECIAL_PKGSUFFIX"}] self.check_referenced("${SRC_URI}", layers) + + +class EnvironmentTests(unittest.TestCase): + def test_environment(self): + os.environ["A"] = "this is A" + self.assertIn("A", os.environ) + self.assertEqual(os.environ["A"], "this is A") + self.assertNotIn("B", os.environ) + + with bb.utils.environment(B="this is B"): + self.assertIn("A", os.environ) + self.assertEqual(os.environ["A"], "this is A") + self.assertIn("B", os.environ) + self.assertEqual(os.environ["B"], "this is B") + + self.assertIn("A", os.environ) + self.assertEqual(os.environ["A"], "this is A") + self.assertNotIn("B", os.environ) diff --git a/lib/bb/utils.py b/lib/bb/utils.py index e6e82d111..70634910f 100644 --- a/lib/bb/utils.py +++ b/lib/bb/utils.py @@ -1681,3 +1681,19 @@ def rename(src, dst): shutil.move(src, dst) else: raise err + +@contextmanager +def environment(**envvars): + """ + Context manager to selectively update the environment with the specified mapping. + """ + backup = dict(os.environ) + try: + os.environ.update(envvars) + yield + finally: + for var in envvars: + if var in backup: + os.environ[var] = backup[var] + else: + del os.environ[var] |