mktempdir: support user-defined prefix

Also switch to using libuv, as it now provides a platform-independent implementation.

Co-Authored-By: Jameson Nash <[email protected]>

Author: sambitdash

NEWS.md

@@ -46,6 +46,7 @@ Standard library changes
 * `filter` now supports `SkipMissing`-wrapped arrays ([#31235]).
 * A no-argument construct to `Ptr{T}` has been added which constructs a null pointer ([#30919])
 * `strip` now accepts a function argument in the same manner as `lstrip` and `rstrip` ([#31211])
+* `mktempdir` now accepts a `prefix` keyword argument to customize the file name ([#31230], [#22922])
 
 #### LinearAlgebra
 

base/file.jl

@@ -416,24 +416,29 @@ function touch(path::AbstractString)
     path
 end
 
+const temp_prefix = "jl_"
+
 if Sys.iswindows()
 
 function tempdir()
     temppath = Vector{UInt16}(undef, 32767)
-    lentemppath = ccall(:GetTempPathW,stdcall,UInt32,(UInt32,Ptr{UInt16}),length(temppath),temppath)
+    lentemppath = ccall(:GetTempPathW, stdcall, UInt32, (UInt32, Ptr{UInt16}), length(temppath), temppath)
     windowserror("GetTempPath", lentemppath >= length(temppath) || lentemppath == 0)
-    resize!(temppath,lentemppath)
+    resize!(temppath, lentemppath)
     return transcode(String, temppath)
 end
 
-const temp_prefix = cwstring("jl_")
 function _win_tempname(temppath::AbstractString, uunique::UInt32)
     tempp = cwstring(temppath)
+    temppfx = cwstring(temp_prefix)
     tname = Vector{UInt16}(undef, 32767)
-    uunique = ccall(:GetTempFileNameW,stdcall,UInt32,(Ptr{UInt16},Ptr{UInt16},UInt32,Ptr{UInt16}), tempp,temp_prefix,uunique,tname)
-    lentname = something(findfirst(iszero,tname), 0)-1
-    windowserror("GetTempFileName", uunique == 0 || lentname <= 0)
-    resize!(tname,lentname)
+    uunique = ccall(:GetTempFileNameW, stdcall, UInt32,
+                    (Ptr{UInt16}, Ptr{UInt16}, UInt32, Ptr{UInt16}),
+                    tempp, temppfx, uunique, tname)
+    windowserror("GetTempFileName", uunique == 0)
+    lentname = something(findfirst(iszero, tname))
+    @assert lentname > 0
+    resize!(tname, lentname - 1)
     return transcode(String, tname)
 end
 
@@ -442,22 +447,6 @@ function mktemp(parent=tempdir())
     return (filename, Base.open(filename, "r+"))
 end
 
-function mktempdir(parent=tempdir())
-    seed::UInt32 = Libc.rand(UInt32)
-    while true
-        if (seed & typemax(UInt16)) == 0
-            seed += 1
-        end
-        filename = _win_tempname(parent, seed)
-        ret = ccall(:_wmkdir, Int32, (Ptr{UInt16},), cwstring(filename))
-        if ret == 0
-            return filename
-        end
-        systemerror(:mktempdir, Libc.errno()!=Libc.EEXIST)
-        seed += 1
-    end
-end
-
 function tempname()
     parent = tempdir()
     seed::UInt32 = rand(UInt32)
@@ -477,7 +466,7 @@ else # !windows
 # Obtain a temporary filename.
 function tempname()
     d = get(ENV, "TMPDIR", C_NULL) # tempnam ignores TMPDIR on darwin
-    p = ccall(:tempnam, Cstring, (Cstring,Cstring), d, :julia)
+    p = ccall(:tempnam, Cstring, (Cstring, Cstring), d, temp_prefix)
     systemerror(:tempnam, p == C_NULL)
     s = unsafe_string(p)
     Libc.free(p)
@@ -489,19 +478,12 @@ tempdir() = dirname(tempname())
 
 # Create and return the name of a temporary file along with an IOStream
 function mktemp(parent=tempdir())
-    b = joinpath(parent, "tmpXXXXXX")
+    b = joinpath(parent, temp_prefix * "XXXXXX")
     p = ccall(:mkstemp, Int32, (Cstring,), b) # modifies b
     systemerror(:mktemp, p == -1)
     return (b, fdio(p, true))
 end
 
-# Create and return the name of a temporary directory
-function mktempdir(parent=tempdir())
-    b = joinpath(parent, "tmpXXXXXX")
-    p = ccall(:mkdtemp, Cstring, (Cstring,), b)
-    systemerror(:mktempdir, p == C_NULL)
-    return unsafe_string(p)
-end
 
 end # os-test
 
@@ -536,12 +518,37 @@ is an open file object for this path.
 mktemp(parent)
 
 """
-    mktempdir(parent=tempdir())
+    mktempdir(parent=tempdir(); prefix=$(repr(temp_prefix)))
 
-Create a temporary directory in the `parent` directory and return its path.
+Create a temporary directory in the `parent` directory with a name
+constructed from the given prefix and a random suffix, and return its path.
+Additionally, any trailing `X` characters may be replaced with random characters.
 If `parent` does not exist, throw an error.
 """
-mktempdir(parent)
+function mktempdir(parent=tempdir(); prefix=temp_prefix)
+    if isempty(parent) || occursin(path_separator_re, parent[end:end])
+        # append a path_separator only if parent didn't already have one
+        tpath = "$(parent)$(prefix)XXXXXX"
+    else
+        tpath = "$(parent)$(path_separator)$(prefix)XXXXXX"
+    end
+
+    req = Libc.malloc(_sizeof_uv_fs)
+    try
+        ret = ccall(:uv_fs_mkdtemp, Int32,
+                    (Ptr{Cvoid}, Ptr{Cvoid}, Cstring, Ptr{Cvoid}),
+                    eventloop(), req, tpath, C_NULL)
+        if ret < 0
+            ccall(:uv_fs_req_cleanup, Cvoid, (Ptr{Cvoid},), req)
+            uv_error("mktempdir", ret)
+        end
+        path = unsafe_string(ccall(:jl_uv_fs_t_path, Cstring, (Ptr{Cvoid},), req))
+        ccall(:uv_fs_req_cleanup, Cvoid, (Ptr{Cvoid},), req)
+        return path
+    finally
+        Libc.free(req)
+    end
+end
 
 
 """
@@ -566,13 +573,13 @@ function mktemp(fn::Function, parent=tempdir())
 end
 
 """
-    mktempdir(f::Function, parent=tempdir())
+    mktempdir(f::Function, parent=tempdir(); prefix=$(repr(temp_prefix)))
 
-Apply the function `f` to the result of [`mktempdir(parent)`](@ref) and remove the
-temporary directory upon completion.
+Apply the function `f` to the result of [`mktempdir(parent; prefix)`](@ref) and remove the
+temporary directory all of its contents upon completion.
 """
-function mktempdir(fn::Function, parent=tempdir())
-    tmpdir = mktempdir(parent)
+function mktempdir(fn::Function, parent=tempdir(); prefix=temp_prefix)
+    tmpdir = mktempdir(parent; prefix=prefix)
     try
         fn(tmpdir)
     finally
@@ -809,7 +816,7 @@ function readlink(path::AbstractString)
             uv_error("readlink", ret)
             @assert false
         end
-        tgt = unsafe_string(ccall(:jl_uv_fs_t_ptr, Ptr{Cchar}, (Ptr{Cvoid},), req))
+        tgt = unsafe_string(ccall(:jl_uv_fs_t_ptr, Cstring, (Ptr{Cvoid},), req))
         ccall(:uv_fs_req_cleanup, Cvoid, (Ptr{Cvoid},), req)
         return tgt
     finally

src/sys.c

@@ -120,6 +120,7 @@ JL_DLLEXPORT int32_t jl_nb_available(ios_t *s)
 JL_DLLEXPORT int jl_sizeof_uv_fs_t(void) { return sizeof(uv_fs_t); }
 JL_DLLEXPORT void jl_uv_fs_req_cleanup(uv_fs_t *req) { uv_fs_req_cleanup(req); }
 JL_DLLEXPORT char *jl_uv_fs_t_ptr(uv_fs_t *req) { return (char*)req->ptr; }
+JL_DLLEXPORT char *jl_uv_fs_t_path(uv_fs_t *req) { return (char*)req->path; }
 JL_DLLEXPORT ssize_t jl_uv_fs_result(uv_fs_t *f) { return f->result; }
 
 // --- stat ---

test/file.jl

@@ -1069,3 +1069,49 @@ let n = tempname()
 end
 
 @test_throws ArgumentError mkpath("fakepath", mode = -1)
+
+@testset "mktempdir 'prefix' argument" begin
+    tmpdirbase = joinpath(tempdir(), "")
+    def_prefix = "jl_"
+    mktempdir() do tmpdir
+        @test isdir(tmpdir)
+        @test startswith(tmpdir, tmpdirbase * def_prefix)
+        @test sizeof(tmpdir) == sizeof(tmpdirbase) + sizeof(def_prefix) + 6
+        @test sizeof(basename(tmpdir)) == sizeof(def_prefix) + 6
+        cd(tmpdir) do
+            Sys.iswindows() || mkdir(".\\")
+            for relpath in (".", "./", ".\\", "")
+                mktempdir(relpath) do tmpdir2
+                    pfx = joinpath(relpath, def_prefix)
+                    @test sizeof(tmpdir2) == sizeof(pfx) + 6
+                    @test startswith(tmpdir2, pfx)
+                end
+            end
+        end
+    end
+    # Special character prefix tests
+    for tst_prefix in ("ABCDEF", "./pfx", ".\\pfx", "", "#!@%^&()-", "/", "\\", "////abc", "\\\\\\\\abc", "∃x∀y")
+        mktempdir(; prefix=tst_prefix) do tmpdir
+            @test isdir(tmpdir)
+            @test startswith(tmpdir, tmpdirbase * tst_prefix)
+            @test sizeof(basename(tmpdir)) == 6 + sizeof(basename(tst_prefix))
+        end
+    end
+
+    @test_throws Base.IOError mktempdir(; prefix="dir_notexisting/bar")
+    @test_throws Base.IOError mktempdir(; prefix="dir_notexisting/")
+    @test_throws Base.IOError mktempdir("dir_notexisting/")
+
+    # Behavioral differences across OS types
+    if Sys.iswindows()
+        # invalid file name
+        @test_throws Base.IOError mktempdir(; prefix="a*b")
+        @test_throws Base.IOError mktempdir("a*b")
+    end
+
+    mktempdir(""; prefix=tmpdirbase) do tmpdir
+        @test startswith(tmpdir, tmpdirbase)
+        @test sizeof(tmpdir) == 6 + sizeof(tmpdirbase)
+        @test sizeof(basename(tmpdir)) == 6
+    end
+end