loopy/ipython_ext.py

-from __future__ import division
+from __future__ import annotations
 
-from IPython.core.magic import (magics_class, Magics, cell_magic)
+from IPython.core.magic import Magics, cell_magic, magics_class
 
 import loopy as lp
 
@@ -9,19 +9,16 @@ import loopy as lp
 class LoopyMagics(Magics):
     @cell_magic
     def fortran_kernel(self, line, cell):
-        result = lp.parse_fortran(cell.encode())
-
-        for knl in result:
-            self.shell.user_ns[knl.name] = knl
+        result = lp.parse_fortran(cell)
+        self.shell.user_ns["prog"] = result
 
     @cell_magic
     def transformed_fortran_kernel(self, line, cell):
         result = lp.parse_transformed_fortran(
-                cell.encode(),
+                cell,
                 transform_code_context=self.shell.user_ns)
 
-        for knl in result:
-            self.shell.user_ns[knl.name] = knl
+        self.shell.user_ns["prog"] = result
 
 
 def load_ipython_extension(ip):