feat(magic): add support for jinja2 templates

singlestore-labs · Dec 9, 2024 · b1fbb1a · b1fbb1a
1 parent e12d028
commit b1fbb1a
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/run_personal/magic.py b/run_personal/magic.py
@@ -8,6 +8,7 @@
 from IPython.core.magic import magics_class
 from IPython.core.magic import needs_local_scope
 from IPython.core.magic import no_var_expand
+from jinja2 import Template
 
 
 @magics_class
@@ -28,11 +29,15 @@ def run_personal(self, line: str, local_ns: Any = None) -> Any:
 
           %run_personal personal_file.ipynb
 
+          %run_personal {{ sample_notebook_name }}
+
         """
-        personal_file = line.strip()
+        template = Template(line.strip())
+        personal_file = template.render(local_ns)
         if not personal_file:
             raise ValueError('No personal file specified.')
-        if personal_file.startswith("'") and personal_file.endswith("'"):
+        if (personal_file.startswith("'") and personal_file.endswith("'")) or \
+           (personal_file.startswith('"') and personal_file.endswith('"')):
             personal_file = personal_file[1:-1]
             if not personal_file:
                 raise ValueError('No personal file specified.')

diff --git a/run_shared/magic.py b/run_shared/magic.py
@@ -8,6 +8,7 @@
 from IPython.core.magic import magics_class
 from IPython.core.magic import needs_local_scope
 from IPython.core.magic import no_var_expand
+from jinja2 import Template
 
 
 @magics_class
@@ -28,11 +29,15 @@ def run_shared(self, line: str, local_ns: Any = None) -> Any:
 
           %run_shared shared_file.ipynb
 
+          %run_shared '{{ sample_notebook_name }}'
+
         """
-        shared_file = line.strip()
+        template = Template(line.strip())
+        shared_file = template.render(local_ns)
         if not shared_file:
             raise ValueError('No shared file specified.')
-        if shared_file.startswith("'") and shared_file.endswith("'"):
+        if (shared_file.startswith("'") and shared_file.endswith("'")) or \
+           (shared_file.startswith('"') and shared_file.endswith('"')):
             shared_file = shared_file[1:-1]
             if not shared_file:
                 raise ValueError('No personal file specified.')