add the ability to remove shortcut

setup_jupyter.py

@@ -10,6 +10,7 @@ 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
@@ -308,7 +309,7 @@ def ensure_env(env_name: str) -> Path:
         return env_prefix            
 
 
-def main(target_env_name: str) -> int:
+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")
         # call conda run to re-run this in the base prefix
@@ -328,6 +329,41 @@ def main(target_env_name: str) -> int:
             )
 
             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
 
@@ -380,6 +416,11 @@ if __name__ == "__main__":
         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"
@@ -401,7 +442,7 @@ if __name__ == "__main__":
     logging.captureWarnings(True)
 
     try:
-        exit(main(args.name))
+        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.