2018-10-16 03:04:35 +03:00
|
|
|
#!/usr/bin/env python3
|
2019-06-20 02:58:25 +03:00
|
|
|
# Copyright (c) Facebook, Inc. and its affiliates.
|
2018-10-16 03:04:35 +03:00
|
|
|
#
|
2019-06-20 02:58:25 +03:00
|
|
|
# This software may be used and distributed according to the terms of the
|
|
|
|
# GNU General Public License version 2.
|
2018-10-16 03:04:35 +03:00
|
|
|
|
|
|
|
import abc
|
|
|
|
import os
|
|
|
|
import typing
|
|
|
|
|
|
|
|
|
|
|
|
class EnvironmentVariableMixin(metaclass=abc.ABCMeta):
|
|
|
|
def set_environment_variable(self, name: str, value: str) -> None:
|
2018-10-17 04:39:13 +03:00
|
|
|
self.__add_cleanup_for_environment_variable(name)
|
|
|
|
os.environ[name] = value
|
|
|
|
|
2018-12-07 02:58:09 +03:00
|
|
|
def set_environment_variables(self, variables: typing.Mapping[str, str]) -> None:
|
|
|
|
for name, value in variables.items():
|
|
|
|
self.set_environment_variable(name, value)
|
|
|
|
|
2018-10-17 04:39:13 +03:00
|
|
|
def unset_environment_variable(self, name: str) -> None:
|
|
|
|
self.__add_cleanup_for_environment_variable(name)
|
2018-12-11 03:27:59 +03:00
|
|
|
self.__unset_environment_variable_with_cleanup(name)
|
2018-10-17 04:39:13 +03:00
|
|
|
|
|
|
|
def __add_cleanup_for_environment_variable(self, name: str) -> None:
|
2018-10-16 03:04:35 +03:00
|
|
|
old_value = os.getenv(name)
|
|
|
|
|
|
|
|
def restore() -> None:
|
|
|
|
if old_value is None:
|
2019-10-17 02:46:42 +03:00
|
|
|
# pyre-fixme[16]: `EnvironmentVariableMixin` has no attribute
|
|
|
|
# `__unset_environment_variable_with_cleanup`.
|
2018-12-11 03:27:59 +03:00
|
|
|
self.__unset_environment_variable_with_cleanup(name)
|
2018-10-16 03:04:35 +03:00
|
|
|
else:
|
|
|
|
os.environ[name] = old_value
|
|
|
|
|
|
|
|
self.addCleanup(restore)
|
|
|
|
|
2018-12-11 03:27:59 +03:00
|
|
|
def __unset_environment_variable_with_cleanup(self, name: str) -> None:
|
|
|
|
try:
|
|
|
|
del os.environ[name]
|
|
|
|
except KeyError:
|
|
|
|
pass
|
|
|
|
|
2018-10-16 03:04:35 +03:00
|
|
|
def addCleanup(
|
|
|
|
self,
|
|
|
|
function: typing.Callable[..., typing.Any],
|
|
|
|
*args: typing.Any,
|
|
|
|
**kwargs: typing.Any
|
|
|
|
) -> None:
|
|
|
|
raise NotImplementedError()
|