Add custom extension support for Tempfile (#5264)

crystal-lang · Nov 10, 2017 · 4b6f1c1 · 4b6f1c1
1 parent b9ea898
commit 4b6f1c1
Show file tree

Hide file tree

Showing 8 changed files with 36 additions and 8 deletions.
diff --git a/spec/std/tempfile_spec.cr b/spec/std/tempfile_spec.cr
@@ -11,6 +11,11 @@ describe Tempfile do
     File.read(tempfile.path).should eq("Hello!")
   end
 
+  it "has given extension if passed to constructor" do
+    tempfile = Tempfile.new "foo", ".pdf"
+    File.extname(tempfile.path).should eq(".pdf")
+  end
+
   it "creates and deletes" do
     tempfile = Tempfile.new "foo"
     tempfile.close
@@ -26,6 +31,11 @@ describe Tempfile do
     File.exists?(tempfile.path).should be_true
   end
 
+  it "has given extension if passed to open" do
+    tempfile = Tempfile.open("foo", ".pdf") { |f| }
+    File.extname(tempfile.path).should eq(".pdf")
+  end
+
   it "creates and writes with TMPDIR environment variable" do
     old_tmpdir = ENV["TMPDIR"]?
     ENV["TMPDIR"] = "/tmp"

diff --git a/src/lib_c/i686-linux-gnu/c/stdlib.cr b/src/lib_c/i686-linux-gnu/c/stdlib.cr
@@ -14,6 +14,7 @@ lib LibC
   fun getenv(name : Char*) : Char*
   fun malloc(size : SizeT) : Void*
   fun mkstemp(template : Char*) : Int
+  fun mkstemps(template : Char*, suffixlen : Int) : Int
   fun putenv(string : Char*) : Int
   fun realloc(ptr : Void*, size : SizeT) : Void*
   fun realpath(name : Char*, resolved : Char*) : Char*

diff --git a/src/lib_c/i686-linux-musl/c/stdlib.cr b/src/lib_c/i686-linux-musl/c/stdlib.cr
@@ -14,6 +14,7 @@ lib LibC
   fun getenv(x0 : Char*) : Char*
   fun malloc(x0 : SizeT) : Void*
   fun mkstemp(x0 : Char*) : Int
+  fun mkstemps(x0 : Char*, x1 : Int) : Int
   fun putenv(x0 : Char*) : Int
   fun realloc(x0 : Void*, x1 : SizeT) : Void*
   fun realpath(x0 : Char*, x1 : Char*) : Char*

diff --git a/src/lib_c/x86_64-linux-gnu/c/stdlib.cr b/src/lib_c/x86_64-linux-gnu/c/stdlib.cr
@@ -14,6 +14,7 @@ lib LibC
   fun getenv(name : Char*) : Char*
   fun malloc(size : SizeT) : Void*
   fun mkstemp(template : Char*) : Int
+  fun mkstemps(template : Char*, suffixlen : Int) : Int
   fun putenv(string : Char*) : Int
   fun realloc(ptr : Void*, size : SizeT) : Void*
   fun realpath(name : Char*, resolved : Char*) : Char*

diff --git a/src/lib_c/x86_64-linux-musl/c/stdlib.cr b/src/lib_c/x86_64-linux-musl/c/stdlib.cr
@@ -14,6 +14,7 @@ lib LibC
   fun getenv(x0 : Char*) : Char*
   fun malloc(x0 : SizeT) : Void*
   fun mkstemp(x0 : Char*) : Int
+  fun mkstemps(x0 : Char*, x1 : Int) : Int
   fun putenv(x0 : Char*) : Int
   fun realloc(x0 : Void*, x1 : SizeT) : Void*
   fun realpath(x0 : Char*, x1 : Char*) : Char*

diff --git a/src/lib_c/x86_64-macosx-darwin/c/stdlib.cr b/src/lib_c/x86_64-macosx-darwin/c/stdlib.cr
@@ -14,6 +14,7 @@ lib LibC
   fun getenv(x0 : Char*) : Char*
   fun malloc(x0 : SizeT) : Void*
   fun mkstemp(x0 : Char*) : Int
+  fun mkstemps(x0 : Char*, x1 : Int) : Int
   fun putenv(x0 : Char*) : Int
   fun realloc(x0 : Void*, x1 : SizeT) : Void*
   fun realpath(x0 : Char*, x1 : Char*) : Char*

diff --git a/src/lib_c/x86_64-portbld-freebsd/c/stdlib.cr b/src/lib_c/x86_64-portbld-freebsd/c/stdlib.cr
@@ -14,6 +14,7 @@ lib LibC
   fun getenv(x0 : Char*) : Char*
   fun malloc(x0 : SizeT) : Void*
   fun mkstemp(x0 : Char*) : Int
+  fun mkstemps(x0 : Char*, x1 : Int) : Int
   fun putenv(x0 : Char*) : Int
   fun realloc(x0 : Void*, x1 : SizeT) : Void*
   fun realpath(x0 : Char*, x1 : Char*) : Char*

diff --git a/src/tempfile.cr b/src/tempfile.cr
@@ -31,12 +31,24 @@ require "c/stdlib"
 # tempfile = Tempfile.new("foo")
 # tempfile.unlink
 # ```
+#
+# The optional `extension` argument can be used to force the resulting filename
+# to end with the given extension.
+#
+# ```
+# Tempfile.new("foo", ".png").path # => "/tmp/foo.ulBCPS.png"
+# ```
 class Tempfile < IO::FileDescriptor
-  # Creates a `Tempfile` with the given filename.
-  def initialize(name)
+  # Creates a `Tempfile` with the given filename and extension.
+  def initialize(name, extension = nil)
     tmpdir = self.class.dirname + File::SEPARATOR
-    @path = "#{tmpdir}#{name}.XXXXXX"
-    fileno = LibC.mkstemp(@path)
+    @path = "#{tmpdir}#{name}.XXXXXX#{extension}"
+    fileno = if extension
+               LibC.mkstemps(@path, extension.bytesize)
+             else
+               LibC.mkstemp(@path)
+             end
+
     if fileno == -1
       raise Errno.new("mkstemp")
     end
@@ -50,17 +62,17 @@ class Tempfile < IO::FileDescriptor
   # ```
   getter path : String
 
-  # Creates a file with *filename*, and yields it to the given block.
-  # It is closed and returned at the end of this method call.
+  # Creates a file with *filename* and *extension*, and yields it to the given
+  # block. It is closed and returned at the end of this method call.
   #
   # ```
   # tempfile = Tempfile.open("foo") do |file|
   #   file.print("bar")
   # end
   # File.read(tempfile.path) # => "bar"
   # ```
-  def self.open(filename)
-    tempfile = Tempfile.new(filename)
+  def self.open(filename, extension = nil)
+    tempfile = Tempfile.new(filename, extension)
     begin
       yield tempfile
     ensure