From 86ad657c1a49b99ce4023f4a1ef48970f96ebb57 Mon Sep 17 00:00:00 2001
From: Andrei Nechaev <anechaev@nuxeo.com>
Date: Wed, 5 Dec 2018 12:14:28 -0500
Subject: [PATCH] supporting multiple files

---
 nvcc_plugin.py | 74 +++++++++++++++++++++++++++++++++++++++++++++++++-
 setup.py       |  2 +-
 2 files changed, 74 insertions(+), 2 deletions(-)

diff --git a/nvcc_plugin.py b/nvcc_plugin.py
index c6dc1b0..880384e 100644
--- a/nvcc_plugin.py
+++ b/nvcc_plugin.py
@@ -5,6 +5,8 @@ import argparse
 import tempfile
 import subprocess
 import IPython.core.magic as ipym
+from IPython.core.magic_arguments import (argument, magic_arguments,
+                                          parse_argstring)
 
 compiler = '/usr/local/cuda/bin/nvcc'
 ext = '.cu'
@@ -30,7 +32,7 @@ class NVCCPlugin(ipym.Magics):
 
     def run(self, file_path, timeit=False):
         if timeit:
-            stmt = "subprocess.check_output(['{}'], stderr=subprocess.STDOUT)".format(file_path + ".out")
+            stmt = f"subprocess.check_output(['{file_path}.out'], stderr=subprocess.STDOUT)"
             output = self.shell.run_cell_magic(magic_name="timeit", line="-q -o import subprocess", cell=stmt)
         else:
             output = subprocess.check_output([file_path + ".out"], stderr=subprocess.STDOUT)
@@ -58,6 +60,76 @@ class NVCCPlugin(ipym.Magics):
         return output
 
 
+out = "result.out"
+
+
+@ipym.magics_class
+class NVCCPluginV2(ipym.Magics):
+
+    def __init__(self, shell):
+        super(NVCCPluginV2, self).__init__(shell)
+        with tempfile.TemporaryDirectory() as tmp:
+            self.output_dir = os.path.join(tmp, str(uuid.uuid4()))
+
+    @staticmethod
+    def compile(file_path):
+        subprocess.check_output([compiler, file_path, "-o", out], stderr=subprocess.STDOUT)
+
+    def run(self, file_path, timeit=False):
+        if timeit:
+            stmt = f"subprocess.check_output(['{out}'], stderr=subprocess.STDOUT)"
+            output = self.shell.run_cell_magic(magic_name="timeit", line="-q -o import subprocess", cell=stmt)
+        else:
+            output = subprocess.check_output([file_path + ".out"], stderr=subprocess.STDOUT)
+            output = output.decode('utf8')
+        return output
+
+    @magic_arguments
+    @argument('-n', '--name', type=str, help='file name that will be produced by the cell. must end with .cu extension')
+    @argument('-c', '--compile', type=bool, help='Should be compiled?')
+    @ipym.cell_magic
+    def cuda(self, line='', cell=None):
+        args = parse_argstring(self.cuda, line)
+        if args.name[:-3] != '.cu':
+            raise Exception('name must end with .cu')
+
+        file_path = os.path.join(self.tmp_dir, args.name)
+        with open(file_path, "w") as f:
+            f.write(cell)
+
+        if args.compile:
+            try:
+                self.compile(file_path)
+                output = self.run(file_path, timeit=args.timeit)
+            except subprocess.CalledProcessError as e:
+                print(e.output.decode("utf8"))
+                output = None
+        else:
+            output = f'File written in {file_path}'
+
+        return output
+
+    @ipym.cell_magic
+    def cuda_run(self, line='', cell=None):
+        try:
+            args = self.argparser.parse_args(line.split())
+        except SystemExit:
+            self.argparser.print_help()
+            return
+
+        try:
+            self.compile('*.cu')
+            output = self.run(out, timeit=args.timeit)
+        except subprocess.CalledProcessError as e:
+            print(e.output.decode("utf8"))
+            output = None
+
+        return output
+
+
 def load_ipython_extension(ip):
     nvcc_plugin = NVCCPlugin(ip)
     ip.register_magics(nvcc_plugin)
+
+    nvcc_plugin_v2 = NVCCPluginV2(ip)
+    ip.register_magics(nvcc_plugin_v2)
diff --git a/setup.py b/setup.py
index 3d3f20e..4bbe10e 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@ from distutils.core import setup
 
 setup(
     name='NVCCPlugin',
-    version='0.0.1',
+    version='0.0.2',
     author='Andrei Nechaev',
     author_email='lyfaradey@yahoo.com',
     py_modules=['nvcc_plugin'],