Jupyter API: improve data module and api interactions

data module can now load its files from both JupyterHub and a download

improve importlib.reload interaction in user-generated notebooks

have the API include itself in tarball generation

Signed-off-by: Lance-Drane <[email protected]>

ipsframework/_jupyter/api_v1.py

@@ -49,14 +49,17 @@ def generate_tar_from_runids(runids: Union[Iterable[int], int]) -> str:
     """
     tarball_name = f'{datetime.datetime.now(datetime.timezone.utc).isoformat().replace(":", "-").replace("+", "_")}__ips_runs'
     tarball = THIS_DIR / f'{tarball_name}.tar.gz'
-    archive = tarfile.open(tarball, 'w:gz')
+    with tarfile.open(tarball, 'w:gz') as archive:
+        # add API files inside the tarball
+        for api_file in THIS_DIR.glob('api_v*.py'):
+            archive.add(api_file, arcname=os.path.join(tarball_name, api_file.name))
 
-    if isinstance(runids, int):
-        runids = [runids]
+        if isinstance(runids, int):
+            runids = [runids]
 
-    for runid in runids:
-        arcname = os.path.join(tarball_name, str(runid), 'data')
-        archive.add(os.path.join(THIS_DIR, str(runid), 'data'), arcname=arcname)
+        # add runids in directory
+        for runid in runids:
+            arcname = os.path.join(tarball_name, str(runid))
+            archive.add(os.path.join(THIS_DIR, str(runid)), arcname=arcname)
 
-    archive.close()
     return str(tarball)

ipsframework/_jupyter/initializer.py

@@ -16,7 +16,6 @@
 
 import re
 import shutil
-from os.path import sep
 from pathlib import Path
 from typing import Optional
 
@@ -37,13 +36,13 @@ def replace_last(source_string: str, old: str, new: str) -> str:
     return f'{head}{new}{tail}'
 
 
-def _initial_data_file_code(dest: str) -> str:
+def _initial_data_file_code() -> str:
     return f"""# This file should be imported by a jupyter notebook or the generated API. DO NOT EDIT UNTIL IPS RUN IS FINALIZED.
 
 import os
+import pathlib
 
-# NOTE: directory should be sim_name plus the run id from the Portal
-{DIRECTORY_VARIABLE_NAME} = '{str(Path(dest).parent / 'data') + sep}'
+{DIRECTORY_VARIABLE_NAME} = str(pathlib.Path(__file__).resolve().parent / 'data') + os.path.sep
 {DATA_VARIABLE_NAME} = {{
 }}
 """
@@ -69,12 +68,17 @@ def initialize_jupyter_notebook(notebook_dest: str, notebook_src: str):
 
     nb['cells'] = [
         # explicitly mark the IPS cell for users inspecting the file, unused programatically
-        nbf.v4.new_markdown_cell('## Next cell generated by IPS Framework'),
+        nbf.v4.new_markdown_cell("""## Next cell generated by IPS Framework
+
+Execute this cell again to use new data during the simulation.
+"""),
         nbf.v4.new_code_cell(f"""
-from {DATA_MODULE_NAME} import {DATA_VARIABLE_NAME}
 import importlib
 
-importlib.reload('{DATA_VARIABLE_NAME}')
+import {DATA_MODULE_NAME}
+importlib.reload({DATA_MODULE_NAME})
+{DATA_VARIABLE_NAME} = {DATA_MODULE_NAME}.{DATA_VARIABLE_NAME}
+
 """),
     ] + nb['cells'][:]
 
@@ -92,7 +96,7 @@ def initialize_jupyter_import_module_file(dest: str):
 
     dest = f'{dest}{DATA_MODULE_NAME}.py'
     with open(dest, 'w') as f:
-        f.write(_initial_data_file_code(dest))
+        f.write(_initial_data_file_code())
 
 
 def update_module_file_with_data_file(dest: str, data_file: str, replace: bool, timestamp: float = 0.0) -> Optional[str]: