[stdlib] first implementaion of NamedtemporaryFile

Signed-off-by: Artemio Garza Reyna <artemiogr97@gmail.com>

stdlib/src/tempfile/__init__.mojo

@@ -0,0 +1,15 @@
+# ===----------------------------------------------------------------------=== #
+# Copyright (c) 2024, Modular Inc. All rights reserved.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions:
+# https://llvm.org/LICENSE.txt
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===----------------------------------------------------------------------=== #
+"""Implements the tempfile package."""
+
+from .tempfile import NamedTemporaryFile

stdlib/src/tempfile/tempfile.mojo

@@ -0,0 +1,233 @@
+# ===----------------------------------------------------------------------=== #
+# Copyright (c) 2024, Modular Inc. All rights reserved.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions:
+# https://llvm.org/LICENSE.txt
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===----------------------------------------------------------------------=== #
+"""Implements tempfile methods.
+
+You can import a method from the `tempfile` package. For example:
+
+```mojo
+from tempfile import NamedTemporaryFile
+```
+"""
+
+from collections import Optional
+import os
+import sys
+import pathlib
+from pathlib import Path
+
+
+fn _get_random_name(size: Int = 8) -> String:
+    var characters = String("abcdefghijklmnopqrstuvwxyz0123456789_")
+    var name = String("")
+    random.seed()
+    for _ in range(size):
+        var rand_index = int(random.random_ui64(0, len(characters) - 1))
+        name += characters[rand_index]
+    return name
+
+
+fn _candidate_tempdir_list() -> List[String]:
+    """Generate a list of candidate temporary directories which
+    _get_default_tempdir will try."""
+
+    var dirlist = List[String]()
+    var possible_env_vars = List("TMPDIR", "TEMP", "TMP")
+    var env_var: String
+    var dirname: String
+
+    # First, try the environment.
+    for env_var in possible_env_vars:
+        dirname = os.getenv(env_var[])
+        if dirname:
+            dirlist.append(dirname)
+
+    # Failing that, try OS-specific locations.
+    if sys.os_is_windows():
+        # TODO handle windows
+        pass
+    else:
+        dirlist.extend(
+            List(String("/tmp"), String("/var/tmp"), String("/usr/tmp"))
+        )
+
+    # As a last resort, the current directory.
+    try:
+        dirlist.append(pathlib.path.cwd())
+    except:
+        pass
+
+    return dirlist
+
+
+fn _get_default_tempdir() raises -> String:
+    """Calculate the default directory to use for temporary files.
+
+    We determine whether or not a candidate temp dir is usable by
+    trying to create and write to a file in that directory.  If this
+    is successful, the test file is deleted. To prevent denial of
+    service, the name of the test file must be randomized."""
+    # TODO In python this function is called exactly one such that the default
+    # tmp dir is the same along the program execution,
+    # since there is not a global scope in mojo yet this is not possible for now
+
+    var dirlist = _candidate_tempdir_list()
+    var dir_name: String
+
+    for dir_name in dirlist:
+        if not os.path.isdir(dir_name[]):
+            continue
+        for _ in range(100):
+            var name = _get_random_name()
+            var filename = dir_name[] + "/" + name
+            if os.path.isfile(filename):
+                continue
+
+            try:
+                var temp_file = FileHandle(filename, "w")
+                temp_file.close()
+                os.remove(filename)
+                return dir_name[]
+            except:
+                break
+    raise Error("No usable temporary directory found")
+
+
+struct NamedTemporaryFile:
+    """A handle to a temporary file."""
+
+    var _file_handle: FileHandle
+    """The underlying file handle."""
+    var _delete: Bool
+    var name: String
+    """Name of the file."""
+
+    fn __init__(
+        inout self,
+        mode: String = "w",
+        suffix: String = "",
+        prefix: String = "tmp",
+        dir: Optional[String] = None,
+        delete: Bool = True,
+    ) raises:
+        """Create a named temporary file.
+        This is a wrapper around a `FileHandle`,
+        os.remove is called in close method if `delete` is True.
+
+        Args:
+            mode: The mode to open the file in (the mode can be "r" or "w").
+            suffix: Suffix to use for the file name.
+            prefix: Prefix to use for the file name.
+            dir: Directory in which the file will be created.
+            delete: Whether the file is deleted on close.
+        """
+        var final_dir: Path
+        if not dir:
+            final_dir = Path(_get_default_tempdir())
+        else:
+            final_dir = Path(dir.value()[])
+
+        self._delete = delete
+
+        var MAX_TRIES = 100
+        for _ in range(MAX_TRIES):
+            var potential_name = final_dir / (
+                prefix + _get_random_name() + suffix
+            )
+            if os.path.exists(potential_name):
+                continue
+            try:
+                self._file_handle = FileHandle(potential_name, mode=mode)
+                # TODO for now this name could be relative,
+                # python implementation expands the path,
+                # but several functions are not yet implemented in mojo
+                # i.e. abspath, normpath
+                self.name = potential_name
+                return
+            except:
+                continue
+        raise Error("Failed to create temporary file")
+
+    @always_inline
+    fn __del__(owned self):
+        """Closes the file handle."""
+        try:
+            self.close()
+        except:
+            pass
+
+    fn close(inout self) raises:
+        """Closes the file handle."""
+        self._file_handle.close()
+        if self._delete:
+            os.remove(self.name)
+
+    fn __moveinit__(inout self, owned existing: Self):
+        """Moves constructor for the file handle.
+
+        Args:
+          existing: The existing file handle.
+        """
+        self._file_handle = existing._file_handle^
+        self._delete = existing._delete
+        self.name = existing.name
+
+    @always_inline
+    fn read(self, size: Int64 = -1) raises -> String:
+        """Reads the data from the file.
+
+        Args:
+            size: Requested number of bytes to read.
+
+        Returns:
+          The contents of the file.
+        """
+        return self._file_handle.read(size)
+
+    fn read_bytes(self, size: Int64 = -1) raises -> List[Int8]:
+        """Read from file buffer until we have `size` characters or we hit EOF.
+        If `size` is negative or omitted, read until EOF.
+
+        Args:
+            size: Requested number of bytes to read.
+
+        Returns:
+          The contents of the file.
+        """
+        return self._file_handle.read_bytes(size)
+
+    fn seek(self, offset: UInt64) raises -> UInt64:
+        """Seeks to the given offset in the file.
+
+        Args:
+            offset: The byte offset to seek to from the start of the file.
+
+        Raises:
+            An error if this file handle is invalid, or if file seek returned a
+            failure.
+
+        Returns:
+            The resulting byte offset from the start of the file.
+        """
+        return self._file_handle.seek(offset)
+
+    fn write(self, data: String) raises:
+        """Write the data to the file.
+
+        Args:
+          data: The data to write to the file.
+        """
+        self._file_handle.write(data)
+
+    fn __enter__(owned self) -> Self:
+        """The function to call when entering the context."""
+        return self^

stdlib/test/tempfile/test_tempfile.mojo

@@ -0,0 +1,67 @@
+# ===----------------------------------------------------------------------=== #
+# Copyright (c) 2024, Modular Inc. All rights reserved.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions:
+# https://llvm.org/LICENSE.txt
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===----------------------------------------------------------------------=== #
+# RUN: %mojo-no-debug %s
+
+import os
+from os.path import exists
+from testing import assert_true, assert_false, assert_equal
+from tempfile import NamedTemporaryFile
+
+
+fn test_named_temporary_file_deletion() raises:
+    var tmp_file: NamedTemporaryFile
+    var file_name: String
+
+    with NamedTemporaryFile(prefix="my_prefix") as my_tmp_file:
+        file_name = my_tmp_file.name
+        assert_true(exists(file_name), "Failed to create file " + file_name)
+        assert_true(file_name.split("/")[-1].startswith("my_prefix"))
+    assert_false(exists(file_name), "Failed to delete file " + file_name)
+
+    with NamedTemporaryFile(delete=False) as my_tmp_file:
+        file_name = my_tmp_file.name
+        assert_true(exists(file_name), "Failed to create file " + file_name)
+    assert_true(exists(file_name), "File " + file_name + " should still exist")
+    os.remove(file_name)
+
+    tmp_file = NamedTemporaryFile()
+    file_name = tmp_file.name
+    assert_true(exists(file_name), "Failed to create file " + file_name)
+    tmp_file.close()
+    assert_false(exists(file_name), "Failed to delete file " + file_name)
+
+    tmp_file = NamedTemporaryFile(delete=False)
+    file_name = tmp_file.name
+    assert_true(exists(file_name), "Failed to create file " + file_name)
+    tmp_file.close()
+    assert_true(exists(file_name), "File " + file_name + " should still exist")
+    os.remove(file_name)
+
+
+fn test_named_temporary_file_write() raises:
+    var file_name: String
+    var contents: String
+
+    with NamedTemporaryFile(delete=False) as my_tmp_file:
+        file_name = my_tmp_file.name
+        my_tmp_file.write("hello world")
+
+    with open(file_name, "r") as my_file:
+        contents = my_file.read()
+    assert_equal("hello world", contents)
+    os.remove(file_name)
+
+
+fn main() raises:
+    test_named_temporary_file_deletion()
+    test_named_temporary_file_write()