Create temp directory locally

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

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)