From 2a447be92e5520928edbf8b70c58e3c9cec707af Mon Sep 17 00:00:00 2001
From: Philip Cho <chohyu01@cs.washington.edu>
Date: Sun, 6 Oct 2019 01:43:48 -0700
Subject: [PATCH] Create temp directory locally

---
 python/treelite/common/util.py | 16 ----------------
 python/treelite/frontend.py    |  5 +++--
 2 files changed, 3 insertions(+), 18 deletions(-)

diff --git a/python/treelite/common/util.py b/python/treelite/common/util.py
index 72aa87ca..a4fbd2a4 100644
--- a/python/treelite/common/util.py
+++ b/python/treelite/common/util.py
@@ -50,22 +50,6 @@ def _load_ver():
 class TreeliteError(Exception):
   """Error thrown by treelite"""
 
-if PY3:
-  # pylint: disable=W0611
-  from tempfile import TemporaryDirectory
-else:
-  import tempfile
-  class TemporaryDirectory():
-    """Context manager for tempfile.mkdtemp()"""
-    # pylint: disable=R0903
-
-    def __enter__(self):
-      self.name = tempfile.mkdtemp()    # pylint: disable=W0201
-      return self.name
-
-    def __exit__(self, exc_type, exc_value, traceback):
-      shutil.rmtree(self.name)
-
 def lineno():
   """Returns line number"""
   return inspect.currentframe().f_back.f_lineno
diff --git a/python/treelite/frontend.py b/python/treelite/frontend.py
index d20671a9..cb9a1870 100644
--- a/python/treelite/frontend.py
+++ b/python/treelite/frontend.py
@@ -5,8 +5,9 @@
 import collections
 import shutil
 import os
+from tempfile import TemporaryDirectory
 from .common.compat import STRING_TYPES
-from .common.util import c_str, TreeliteError, TemporaryDirectory
+from .common.util import c_str, TreeliteError
 from .core import _LIB, c_array, _check_call
 from .contrib import create_shared, generate_makefile, _toolchain_exist_check
 
@@ -124,7 +125,7 @@ def export_lib(self, toolchain, libpath, params=None, compiler='ast_native',
        shutil.move('/temporary/directory/mymodel.dll', './mymodel.dll')
     """
     _toolchain_exist_check(toolchain)
-    with TemporaryDirectory() as temp_dir:
+    with TemporaryDirectory(dir=os.path.dirname(libpath)) as temp_dir:
       self.compile(temp_dir, params, compiler, verbose)
       temp_libpath = create_shared(toolchain, temp_dir, nthread,
                                    verbose, options)