From b1fbb1ad83f8e59ca30d830b2506b734afd2540e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nuno=20Gon=C3=A7alves?= <ngoncalves@singlestore.com>
Date: Mon, 9 Dec 2024 14:23:08 +0000
Subject: [PATCH] feat(magic): add support for jinja2 templates

---
 run_personal/magic.py | 9 +++++++--
 run_shared/magic.py   | 9 +++++++--
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/run_personal/magic.py b/run_personal/magic.py
index e5dc6bfa..20999974 100644
--- 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
index 87949bde..384797a7 100644
--- 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.')