storage: Add Directory.open_file() method

This is the virtual directory API equivalent to Python's built-in
`open()` function.

3 files changed, 53 insertions, 0 deletions

diff --git a/src/buildstream/storage/_casbaseddirectory.py b/src/buildstream/storage/_casbaseddirectory.py
index ffcf085bb..5f30037af 100644
--- a/src/buildstream/storage/_casbaseddirectory.py
+++ b/src/buildstream/storage/_casbaseddirectory.py
@@ -31,8 +31,10 @@ import os
 import stat
 import copy
 import tarfile as tarfilelib
+from contextlib import contextmanager
 from io import StringIO
 
+from .. import utils
 from .._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
 from .directory import Directory, VirtualDirectoryError, _FileType
 from ._filebaseddirectory import FileBasedDirectory
@@ -731,6 +733,34 @@ class CasBasedDirectory(Directory):
             path += "/" + self.common_name
         return path
 
+    @contextmanager
+    def open_file(self, *path: str, mode: str = "r"):
+        subdir = self.descend(*path[:-1])
+        entry = subdir.index.get(path[-1])
+
+        if entry and entry.type != _FileType.REGULAR_FILE:
+            raise VirtualDirectoryError("{} in {} is not a file".format(path[-1], str(subdir)))
+
+        if mode not in ["r", "rb", "w", "wb", "x", "xb"]:
+            raise ValueError("Unsupported mode: `{}`".format(mode))
+
+        if "r" in mode:
+            if not entry:
+                raise FileNotFoundError("{} not found in {}".format(path[-1], str(subdir)))
+
+            # Read-only access, allow direct access to CAS object
+            with open(self.cas_cache.objpath(entry.digest), mode, encoding="utf-8") as f:
+                yield f
+        else:
+            if "x" in mode and entry:
+                raise FileExistsError("{} already exists in {}".format(path[-1], str(subdir)))
+
+            with utils._tempnamedfile(mode, encoding="utf-8", dir=self.cas_cache.tmpdir) as f:
+                yield f
+                # Import written temporary file into CAS
+                f.flush()
+                subdir._add_file(path[-1], f.name, modified=True)
+
     def __str__(self):
         return "[CAS:{}]".format(self._get_identifier())
 
diff --git a/src/buildstream/storage/_filebaseddirectory.py b/src/buildstream/storage/_filebaseddirectory.py
index 62fcc39c0..3d3116411 100644
--- a/src/buildstream/storage/_filebaseddirectory.py
+++ b/src/buildstream/storage/_filebaseddirectory.py
@@ -248,6 +248,19 @@ class FileBasedDirectory(Directory):
         except (VirtualDirectoryError, FileNotFoundError):
             return False
 
+    def open_file(self, *path: str, mode: str = "r"):
+        # Use descend() to avoid following symlinks (potentially escaping the sandbox)
+        subdir = self.descend(*path[:-1])
+        newpath = os.path.join(subdir.external_directory, path[-1])
+
+        if mode not in ["r", "rb", "w", "wb", "x", "xb"]:
+            raise ValueError("Unsupported mode: `{}`".format(mode))
+
+        if "r" in mode:
+            return open(newpath, mode=mode, encoding="utf-8")
+        else:
+            return utils.save_file_atomic(newpath, mode=mode, encoding="utf-8")
+
     def __str__(self):
         # This returns the whole path (since we don't know where the directory started)
         # which exposes the sandbox directory; we will have to assume for the time being
diff --git a/src/buildstream/storage/directory.py b/src/buildstream/storage/directory.py
index a9249baee..9ce0cacf3 100644
--- a/src/buildstream/storage/directory.py
+++ b/src/buildstream/storage/directory.py
@@ -208,6 +208,16 @@ class Directory:
         """
         raise NotImplementedError()
 
+    def open_file(self, *paths: str, mode: str = "r"):
+        """ Open file and return a corresponding file object. In text mode,
+        UTF-8 is used as encoding.
+
+        Args:
+          *paths: A list of strings which are all path components.
+          mode (str): An optional string that specifies the mode in which the file is opened.
+        """
+        raise NotImplementedError()
+
 
 # FileType:
 #