Add context manager to deactivate env vars
@@ -10,23 +10,25 @@ import sys
import time
log = logging.getLogger(__name__)
+OFF = ["CTADATA", "GAMMA_CAT", "GAMMAPY_EXTRA", "GAMMAPY_FERMI_LAT_DATA"]
@click.command(name="run")
@click.pass_context
-@click.option("--tutoenv", is_flag=True, default=False, help="Tutorials environment?", show_default=True)
-def cli_jupyter_run(ctx, tutoenv):
+@click.option(
+ "--tutor",
+ is_flag=True,
+ default=False,
+ help="Tutorials environment?",
+ show_default=True,
+)
+def cli_jupyter_run(ctx, tutor):
"""Execute Jupyter notebooks."""
- if tutoenv:
- envpths = unsetenv()
-
+ with environment(OFF, tutor):
for path in ctx.obj["paths"]:
execute_notebook(path)
- if tutoenv:
- envback(envpths)
-
def execute_notebook(path, loglevel=20):
"""Execute a Jupyter notebook."""
@@ -129,19 +131,20 @@ class BlackNotebook:
@click.command(name="test")
@click.pass_context
-@click.option("--tutoenv", is_flag=True, default=False, help="Tutorials environment?", show_default=True)
-def cli_jupyter_test(ctx, tutoenv):
+@click.option(
+ "--tutor",
+ is_flag=True,
+ default=False,
+ help="Tutorials environment?",
+ show_default=True,
+)
+def cli_jupyter_test(ctx, tutor):
"""Check if Jupyter notebooks are broken."""
- if tutoenv:
- envpths = unsetenv()
-
+ with environment(OFF, tutor):
for path in ctx.obj["paths"]:
notebook_test(path)
- if tutoenv:
- envback(envpths)
-
def notebook_test(path):
"""Execute and parse a Jupyter notebook exposing broken cells."""
@@ -178,20 +181,25 @@ def notebook_test(path):
return False
-def unsetenv():
+class environment:
+ """
+ Helper for setting environmental variables
+ """
+
+ def __init__(self, envs, tutor):
+ self.envs = envs
+ self.tutor = tutor
+
+ def __enter__(self):
- envpths = {}
- envars = ["GAMMAPY_EXTRA", "CTADATA", "GAMMA_CAT", "GAMMAPY_FERMI_LAT_DATA"]
- for item in envars:
+ self.old = os.environ
+ if self.tutor:
+ for item in self.envs:
if item in os.environ:
- envpths[item] = os.environ[item]
del os.environ[item]
logging.info("Unsetting {} environment variable.".format(item))
- return envpths
-
-
-def envback(envpths):
- for item in envpths:
- os.environ[item] = envpths[item]
- logging.info("Setting {}={}".format(item, envpths[item]))
+ def __exit__(self, type, value, traceback):
+ if self.tutor:
+ os.environ = self.old
+ logging.info("Environment variables recovered.")
