sapling/eden/test_support/environment_variable.py
generatedunixname89002005307016 4c76d686a8 Update pyre version for eden
Summary: Automatic upgrade to remove `version` override and silence errors.

Reviewed By: grievejia

Differential Revision: D17956249

fbshipit-source-id: d5c8b5aa73151b3fea67aec35d70f332030da2c9
2019-10-16 16:56:29 -07:00

51 lines
1.6 KiB
Python

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This software may be used and distributed according to the terms of the
# GNU General Public License version 2.
import abc
import os
import typing
class EnvironmentVariableMixin(metaclass=abc.ABCMeta):
def set_environment_variable(self, name: str, value: str) -> None:
self.__add_cleanup_for_environment_variable(name)
os.environ[name] = value
def set_environment_variables(self, variables: typing.Mapping[str, str]) -> None:
for name, value in variables.items():
self.set_environment_variable(name, value)
def unset_environment_variable(self, name: str) -> None:
self.__add_cleanup_for_environment_variable(name)
self.__unset_environment_variable_with_cleanup(name)
def __add_cleanup_for_environment_variable(self, name: str) -> None:
old_value = os.getenv(name)
def restore() -> None:
if old_value is None:
# pyre-fixme[16]: `EnvironmentVariableMixin` has no attribute
# `__unset_environment_variable_with_cleanup`.
self.__unset_environment_variable_with_cleanup(name)
else:
os.environ[name] = old_value
self.addCleanup(restore)
def __unset_environment_variable_with_cleanup(self, name: str) -> None:
try:
del os.environ[name]
except KeyError:
pass
def addCleanup(
self,
function: typing.Callable[..., typing.Any],
*args: typing.Any,
**kwargs: typing.Any
) -> None:
raise NotImplementedError()