470 lines
16 KiB
Python
Executable File
470 lines
16 KiB
Python
Executable File
#!/bin/env python
|
|
import logging
|
|
import platform
|
|
|
|
from argparse import ArgumentParser
|
|
from json import JSONDecodeError, loads
|
|
from os import PathLike, environ, linesep
|
|
from pathlib import Path
|
|
from shutil import copy
|
|
from subprocess import CalledProcessError, run
|
|
from sys import argv, stderr, stdout
|
|
from tempfile import TemporaryDirectory
|
|
from typing import Union
|
|
from urllib.parse import urlparse
|
|
|
|
from requests import get
|
|
from requests.exceptions import RequestException
|
|
|
|
LOGGER = logging.getLogger(__name__)
|
|
|
|
class OperationCancelled(Exception):
|
|
pass
|
|
|
|
|
|
def get_current_prefix() -> Path:
|
|
try:
|
|
current_prefix = Path(environ["CONDA_PREFIX"])
|
|
except KeyError:
|
|
try:
|
|
# I vaguely recall MacOS prefixing the conda environment variables with _ for some stupid reason
|
|
current_prefix = Path(environ["_CONDA_PREFIX"])
|
|
except KeyError as e:
|
|
raise RuntimeError("No active conda environment detected. Please activate your base conda environment") from e
|
|
|
|
LOGGER.debug(f"Current prefix is {current_prefix}")
|
|
|
|
return current_prefix
|
|
|
|
|
|
def get_base_prefix() -> Path:
|
|
try:
|
|
conda_exe = Path(environ["CONDA_EXE"])
|
|
except KeyError:
|
|
try:
|
|
conda_exe = Path(environ["_CONDA_EXE"])
|
|
except KeyError as e:
|
|
raise RuntimeError("No active conda environment detected. Please activate your base conda environment") from e
|
|
|
|
LOGGER.debug(f"Base prefix is {conda_exe.parents[1]}")
|
|
|
|
# your base environment will be the grandparent of the conda executable
|
|
return conda_exe.parents[1]
|
|
|
|
|
|
def get_menuinst_version() -> tuple[str, str, str]:
|
|
conda_process = run(["conda", "list", "--prefix", str(get_base_prefix()), "--json"], capture_output=True, check=True)
|
|
|
|
try:
|
|
conda_pkgs = loads(conda_process.stdout)
|
|
except JSONDecodeError as decode_exception:
|
|
raise JSONDecodeError(f"Error parsing conda output as JSON: {decode_exception}") from decode_exception
|
|
|
|
for pkg in conda_pkgs:
|
|
if pkg["name"] == "menuinst":
|
|
break
|
|
else:
|
|
raise RuntimeError("menuinst was not found in the base conda envirionment")
|
|
|
|
LOGGER.debug(f"menuinst=={pkg['version']} found")
|
|
|
|
return pkg["version"].split(".")
|
|
|
|
|
|
def in_base_env() -> bool:
|
|
return get_current_prefix().samefile(get_base_prefix())
|
|
|
|
|
|
def menuinst_gt_v2_present() -> bool:
|
|
major, minor, patch = get_menuinst_version()
|
|
|
|
return (int(major) >= 2) and patch.isnumeric()
|
|
|
|
|
|
def meets_prerequisites() -> bool:
|
|
return in_base_env() and menuinst_gt_v2_present()
|
|
|
|
|
|
def stage_configs(destination_dir: Path) -> Path:
|
|
repo_dir = Path(__file__).parent
|
|
shortcut_json = repo_dir / "jupyterlab_shortcut.json"
|
|
jupyterlab_config = repo_dir / "jupyter_lab_config.py"
|
|
|
|
LOGGER.debug(f"Repository directory is {repo_dir}")
|
|
LOGGER.debug(f"menuinst spec json file is {shortcut_json}")
|
|
LOGGER.debug(f"JupyterLab config file is {jupyterlab_config}")
|
|
# download an icon file referenced in shortcut spec json
|
|
download_icon_file(destination_dir)
|
|
|
|
if not destination_dir.samefile(repo_dir):
|
|
LOGGER.debug(f"The config files will be staged in {destination_dir}")
|
|
copy(jupyterlab_config, destination_dir)
|
|
copy(shortcut_json, destination_dir)
|
|
return destination_dir / shortcut_json.name
|
|
else:
|
|
LOGGER.debug(f"Repository directory {repo_dir} is already the environment Menu directory. No need to stage config files")
|
|
return shortcut_json
|
|
|
|
|
|
def download_icon_file(menu_dir: PathLike):
|
|
LOGGER.debug(f"OS is {platform.system()}")
|
|
match platform.system():
|
|
case "Windows":
|
|
url = "https://raw.githubusercontent.com/jupyterlab/jupyterlab-desktop/master/dist-resources/icons/icon.ico"
|
|
case "Linux":
|
|
url = "https://raw.githubusercontent.com/jupyterlab/jupyterlab-desktop/master/dist-resources/icons/512x512.png"
|
|
case "Darwin":
|
|
url = "https://raw.githubusercontent.com/jupyterlab/jupyterlab-desktop/8614277f274b0e9ee9cc550d194e4f02d0c5c3c7/dist-resources/icon.svg"
|
|
case _:
|
|
raise RuntimeError(f"{platform.system()} is not a supported platform")
|
|
|
|
try:
|
|
LOGGER.debug(f"Downloading {url} to {menu_dir} for use as shortcut icon")
|
|
response = get(url)
|
|
except RequestException as e:
|
|
raise RuntimeError("Unable to download an icon to use for the JupyterLab shortcut") from e
|
|
else:
|
|
file_extension = Path(urlparse(url).path).suffix
|
|
with (Path(menu_dir) / f"jupyterlab{file_extension}").open("wb") as icon_file:
|
|
LOGGER.debug(f"Writing icon file to {icon_file.name}")
|
|
icon_file.write(response.content)
|
|
|
|
if platform.system() == "Darwin" and file_extension == ".svg":
|
|
svg_to_icns(menu_dir/"jupyterlab.svg")
|
|
|
|
|
|
def svg_to_icns(svg_file: PathLike) -> None:
|
|
if not platform.system() == "Darwin":
|
|
raise RuntimeError("svg_to_icns is designed to run on MacOS only")
|
|
|
|
svg_file = Path(svg_file)
|
|
|
|
if not svg_file.exists():
|
|
raise ValueError(f"{svg_file} does not exist")
|
|
elif not svg_file.is_file():
|
|
raise ValueError(f"{svg_file} exists but is not a file")
|
|
elif not svg_file.suffix == ".svg":
|
|
raise ValueError(f"{svg_file} doesn't seem to be an SVG file")
|
|
|
|
target_dir = svg_file.parent
|
|
|
|
LOGGER.debug(f"Converting {svg_file} to {svg_file.with_suffix('.icns')}")
|
|
|
|
with TemporaryDirectory() as temp_dir:
|
|
iconset_dir = Path(temp_dir)/"jupyterlab.iconset"
|
|
iconset_dir.mkdir(exist_ok=True)
|
|
|
|
for size in [2**n for n in range(4, 11)]:
|
|
outfile = iconset_dir/f"icon_{size}x{size}.png"
|
|
try:
|
|
run(
|
|
[
|
|
"qlmanage",
|
|
"-t",
|
|
"-s",
|
|
str(size),
|
|
"-o",
|
|
str(iconset_dir),
|
|
str(svg_file)
|
|
],
|
|
capture_output=True,
|
|
check=True
|
|
)
|
|
except CalledProcessError as e:
|
|
LOGGER.critical(f"Error converting SVG icon to PNG thumbnail iconset with qlmanage")
|
|
raise
|
|
|
|
(iconset_dir/svg_file.with_suffix(".svg.png").name).rename(outfile)
|
|
|
|
if size > 16:
|
|
copy(outfile, outfile.with_stem(f"icon_{size/2}x{size/2}@2x"))
|
|
|
|
try:
|
|
run(
|
|
[
|
|
"iconutil",
|
|
"-c",
|
|
"icns",
|
|
"-o",
|
|
str(target_dir/"jupyterlab.icns"),
|
|
str(iconset_dir)
|
|
],
|
|
capture_output=True,
|
|
check=True
|
|
)
|
|
except CalledProcessError as e:
|
|
LOGGER.critical(f"Error calling iconutil to create MacOS .icns file from iconset in {iconset_dir}")
|
|
raise
|
|
|
|
svg_file.unlink()
|
|
|
|
|
|
def ensure_env(env_name: str) -> Path:
|
|
conda_proc = run(
|
|
[
|
|
"conda",
|
|
"list",
|
|
"--json",
|
|
"--name",
|
|
env_name
|
|
],
|
|
capture_output=True,
|
|
check=False
|
|
)
|
|
|
|
try:
|
|
env_spec = loads(conda_proc.stdout)
|
|
except JSONDecodeError as e:
|
|
raise RuntimeError(f"Error parsing conda output while checking the existance of '{env_name}'") from e
|
|
else:
|
|
env_prefix = get_base_prefix()/f"envs/{env_name}"
|
|
|
|
if isinstance(env_spec, dict) and "error" in env_spec.keys():
|
|
# enviornment doesn't exist. create one with some commonly used packages
|
|
# including ipywidgets because it will automatically install the plugin to render them in jupyterlab
|
|
# ultimately, though, nb_conda_kernels allows the execution of notebooks in any installed conda environment
|
|
# which has `ipykernel` installed
|
|
LOGGER.info(f"Creating new conda enviornment '{env_name}'")
|
|
run(
|
|
[
|
|
"conda",
|
|
"create",
|
|
"--yes",
|
|
"--name",
|
|
env_name,
|
|
"--override-channels",
|
|
"--channel",
|
|
"conda-forge",
|
|
"--no-default-packages",
|
|
"jupyterlab",
|
|
"nb_conda_kernels",
|
|
"nbconvert",
|
|
"jupyterlab-nbconvert-nocode",
|
|
"jupyterlab-git",
|
|
"jupyterlab-lsp",
|
|
"jupyterlab_code_formatter",
|
|
"jupyterlab-day",
|
|
"jupyterlab-night",
|
|
"jupyterlab_pygments",
|
|
"black",
|
|
"isort",
|
|
"python-lsp-server",
|
|
"ipykernel",
|
|
"ipywidgets",
|
|
"panel",
|
|
"pandas",
|
|
"numpy",
|
|
"scipy",
|
|
"statsmodels",
|
|
"openpyxl",
|
|
"tabulate",
|
|
"matplotlib",
|
|
"toolz",
|
|
"more-itertools",
|
|
"ipyparallel",
|
|
"requests",
|
|
"tqdm",
|
|
"rich",
|
|
"rich-with-jupyter",
|
|
"sqlalchemy",
|
|
],
|
|
capture_output=False,
|
|
check=True
|
|
)
|
|
|
|
# set conda-forge as the one and only channel for this environment
|
|
with open(env_prefix/".condarc", "w") as condarc:
|
|
print(
|
|
"pip_interop_enabled: true",
|
|
"channels: [conda-forge]",
|
|
sep=linesep,
|
|
end=linesep,
|
|
file=condarc
|
|
)
|
|
|
|
else:
|
|
# add minimal new jupyter packages t existing environment
|
|
pkgs = {pkg["name"] for pkg in env_spec}
|
|
|
|
missing_pkgs = {"jupyterlab", "nb_conda_kernels", "ipykernel"} - pkgs
|
|
|
|
if missing_pkgs:
|
|
LOGGER.info(f"Installing {'and '.join(missing_pkgs)} in '{env_name}")
|
|
|
|
run(
|
|
[
|
|
"conda",
|
|
"install",
|
|
"--yes",
|
|
"--name",
|
|
env_name,
|
|
*sorted(missing_pkgs)
|
|
],
|
|
capture_output=False,
|
|
check=True
|
|
)
|
|
|
|
(env_prefix/"Menu").mkdir(exist_ok=True)
|
|
|
|
return env_prefix
|
|
|
|
|
|
def main(target_env_name: str, remove_shortcut: bool=False) -> Union[int, None]:
|
|
if not in_base_env():
|
|
LOGGER.warning("Not in base environment. Re-running this script from the base environment")
|
|
|
|
# I don't know why, but Windows *really* does not want to resolve conda in PATH if you aren't in the base environment
|
|
try:
|
|
conda_exe = environ["CONDA_EXE"]
|
|
except KeyError:
|
|
try:
|
|
conda_exe = environ["_CONDA_EXE"]
|
|
except KeyError:
|
|
conda_exe = "conda"
|
|
|
|
# call conda run to re-run this in the base prefix
|
|
rerun_proces = run(
|
|
[conda_exe, "run", "--prefix", str(get_base_prefix()), "--no-capture-output", "python", *argv],
|
|
capture_output=False,
|
|
check=False
|
|
)
|
|
# Surface subprocess return code to parent process
|
|
return rerun_proces.returncode
|
|
elif not menuinst_gt_v2_present():
|
|
if input("This script requires menuinst>=2.0.0. Would you like to install it and try again? ").lower() in ("y", "yes"):
|
|
run(
|
|
["conda", "install", "--prefix", str(get_base_prefix()), "menuinst>=2.0.0"],
|
|
capture_output=False,
|
|
check=True
|
|
)
|
|
|
|
run(["python", *argv], capture_output=False, check=True)
|
|
elif meets_prerequisites() and remove_shortcut:
|
|
from menuinst.api import remove
|
|
|
|
target_prefix = get_base_prefix() / f"envs/{target_env_name}"
|
|
menu_dir = target_prefix / "Menu"
|
|
|
|
shortcut_json = menu_dir / "jupyterlab_shortcut.json"
|
|
jupyterlab_config = menu_dir / "jupyter_lab_config.py"
|
|
match platform.system():
|
|
case "Windows":
|
|
icon_file = menu_dir/"jupyterlab.ico"
|
|
case "Linux":
|
|
icon_file = menu_dir/"jupyterlab.png"
|
|
case "Darwin":
|
|
icon_file = menu_dir/"jupyterlab.ico"
|
|
case _:
|
|
raise RuntimeError(f"{platform.system()} is not a supported platform")
|
|
|
|
if shortcut_json.is_file():
|
|
LOGGER.debug("Removing JupyterLab shortcut")
|
|
try:
|
|
remove(shortcut_json, target_prefix=target_prefix)
|
|
except:
|
|
pass
|
|
else:
|
|
LOGGER.info("JupyterLab shortcut removed")
|
|
else:
|
|
LOGGER.error(f"Shortcut spec file '{shortcut_json}' does not exist, therefore the shotcut cannot be removed by this script. Please delete it manually")
|
|
|
|
if menu_dir.exists() and menu_dir.is_dir():
|
|
LOGGER.debug(f"Cleaning up {target_env_name} environment's Menu directory {menu_dir}")
|
|
jupyterlab_config.unlink(missing_ok=True)
|
|
icon_file.unlink(missing_ok=True)
|
|
shortcut_json.unlink(missing_ok=True)
|
|
|
|
elif meets_prerequisites():
|
|
from menuinst.api import install, remove
|
|
|
|
target_prefix = ensure_env(target_env_name)
|
|
|
|
shortcut_json = stage_configs(target_prefix/"Menu")
|
|
|
|
try:
|
|
LOGGER.debug("Removing old shortcut")
|
|
remove(shortcut_json, target_prefix=target_prefix)
|
|
except:
|
|
pass
|
|
finally:
|
|
LOGGER.debug("Creating new shortcut")
|
|
install(shortcut_json, target_prefix=target_prefix)
|
|
LOGGER.info("Jupyter Lab shortcut created!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
LOGGER = logging.getLogger()
|
|
|
|
STDOUT = logging.StreamHandler(stdout)
|
|
STDOUT.setLevel(logging.INFO)
|
|
STDOUT.addFilter(lambda r: r.levelno < logging.WARNING)
|
|
STDOUT.setFormatter(logging.Formatter())
|
|
|
|
STDERR = logging.StreamHandler(stderr)
|
|
STDERR.setLevel(logging.WARNING)
|
|
STDERR.setFormatter(
|
|
logging.Formatter(
|
|
"{levelname}: {message}",
|
|
style="{"
|
|
)
|
|
)
|
|
|
|
DEBUG_HANDLER = logging.StreamHandler(stderr)
|
|
DEBUG_HANDLER.setLevel(logging.DEBUG)
|
|
DEBUG_HANDLER.addFilter(lambda r: r.levelno < logging.INFO)
|
|
DEBUG_HANDLER.setFormatter(
|
|
logging.Formatter(
|
|
"{levelname}: {message}",
|
|
style="{"
|
|
)
|
|
)
|
|
|
|
parser = ArgumentParser()
|
|
|
|
parser.add_argument(
|
|
"--debug",
|
|
action = "store_true",
|
|
help = "Enable debug logging to the terminal"
|
|
)
|
|
parser.add_argument(
|
|
"--remove",
|
|
action = "store_true",
|
|
help = "Remove the shortcut created by this script and cleanup environment Menu directory to remove files created by this script"
|
|
)
|
|
parser.add_argument(
|
|
"name",
|
|
help = "Name of the target conda environment containing JupyterLab. If it does not exist, one will be created"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.debug:
|
|
logging.basicConfig(
|
|
level=logging.DEBUG,
|
|
handlers=(STDOUT, STDERR, DEBUG_HANDLER),
|
|
)
|
|
else:
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
handlers=(STDOUT, STDERR),
|
|
)
|
|
|
|
logging.captureWarnings(True)
|
|
|
|
try:
|
|
logging.info("Creating Jupyterlab shortcut. This may take a moment...")
|
|
exit(main(args.name, args.remove))
|
|
except CalledProcessError as e:
|
|
# catch any exceptions raised if calling conda returns an unsuccessful return code
|
|
# log exception and suface conda's return code.
|
|
logging.exception(f"'{e.cmd}' failed with exit code {e.returncode}:{linesep * 2}{e. output}")
|
|
exit(e.returncode)
|
|
except (OperationCancelled, KeyboardInterrupt) as e:
|
|
logging.info("Operation cancelled by the user")
|
|
exit(1)
|
|
except Exception as e:
|
|
logging.exception(str(e))
|
|
exit(1)
|
|
finally:
|
|
logging.shutdown()
|