diff options
Diffstat (limited to 'tensorflow/python/lib/io/file_io.py')
| -rw-r--r-- | tensorflow/python/lib/io/file_io.py | 65 |
1 files changed, 45 insertions, 20 deletions
diff --git a/tensorflow/python/lib/io/file_io.py b/tensorflow/python/lib/io/file_io.py index f5b138c9da..7a19ae2b71 100644 --- a/tensorflow/python/lib/io/file_io.py +++ b/tensorflow/python/lib/io/file_io.py @@ -45,6 +45,7 @@ class FileIO(object): self.__name = name self.__mode = mode self._read_buf = None + self._writable_file = None if mode not in ("r", "w", "a", "r+", "w+", "a+"): raise errors.InvalidArgumentError( None, None, "mode is not 'r' or 'w' or 'a' or 'r+' or 'w+' or 'a+'") @@ -72,18 +73,27 @@ class FileIO(object): raise errors.InternalError(None, None, "Could not open file for streaming") + def _prewrite_check(self): + if not self._writable_file: + if not self._write_check_passed: + raise errors.PermissionDeniedError(None, None, + "File isn't open for writing") + self._writable_file = pywrap_tensorflow.CreateWritableFile( + compat.as_bytes(self.__name)) + if not self._writable_file: + raise errors.InternalError(None, None, + "Could not open file for writing") + def size(self): """Returns the size of the file.""" return stat(self.__name).length def write(self, file_content): - """Writes file_content to the file.""" - if not self._write_check_passed: - raise errors.PermissionDeniedError(None, None, - "File isn't open for writing") + """Writes file_content to the file. Appends to the end of the file.""" + self._prewrite_check() with errors.raise_exception_on_not_ok_status() as status: - pywrap_tensorflow.WriteStringToFile( - compat.as_bytes(self.__name), compat.as_bytes(file_content), status) + pywrap_tensorflow.AppendToFile( + compat.as_bytes(file_content), self._writable_file, status) def read(self): """Returns the contents of a file as a string.""" @@ -97,7 +107,7 @@ class FileIO(object): def readline(self): r"""Reads the next line from the file. Leaves the '\n' at the end.""" self._prereadline_check() - return self._read_buf.ReadLineAsString() + return compat.as_str_any(self._read_buf.ReadLineAsString()) def readlines(self): """Returns all lines from the file in a list.""" @@ -116,7 +126,7 @@ class FileIO(object): def __exit__(self, unused_type, unused_value, unused_traceback): """Make usable with "with" statement.""" - self._read_buf = None + self.close() def __iter__(self): return self @@ -130,6 +140,18 @@ class FileIO(object): def __next__(self): return self.next() + def flush(self): + """Flushes the Writable file to disk.""" + if self._writable_file: + with errors.raise_exception_on_not_ok_status() as status: + pywrap_tensorflow.FlushWritableFile(self._writable_file, status) + + def close(self): + """Closes FileIO. Should be called for the WritableFile to be flushed.""" + self._read_buf = None + self.flush() + self._writable_file = None + def file_exists(filename): """Determines whether a path exists or not. @@ -184,8 +206,8 @@ def write_string_to_file(filename, file_content): Raises: errors.OpError: If there are errors during the operation. """ - f = FileIO(filename, mode="w") - f.write(file_content) + with FileIO(filename, mode="w") as f: + f.write(file_content) def get_matching_files(filename): @@ -203,7 +225,7 @@ def get_matching_files(filename): with errors.raise_exception_on_not_ok_status() as status: # Convert each element to string, since the return values of the # vector of string should be interpreted as strings, not bytes. - return [compat.as_str(matching_filename) + return [compat.as_str_any(matching_filename) for matching_filename in pywrap_tensorflow.GetMatchingFiles( compat.as_bytes(filename), status)] @@ -235,7 +257,7 @@ def recursive_create_dir(dirname): errors.OpError: If the operation fails. """ with errors.raise_exception_on_not_ok_status() as status: - dirs = dirname.split("/") + dirs = compat.as_str_any(dirname).split("/") for i in range(len(dirs)): partial_dir = "/".join(dirs[0:i + 1]) if partial_dir and not file_exists(partial_dir): @@ -301,8 +323,8 @@ def is_directory(dirname): Raises: errors.OpError: If the path doesn't exist or other errors """ - with errors.raise_exception_on_not_ok_status() as status: - return pywrap_tensorflow.IsDirectory(compat.as_bytes(dirname), status) + status = pywrap_tensorflow.TF_NewStatus() + return pywrap_tensorflow.IsDirectory(compat.as_bytes(dirname), status) def list_directory(dirname): @@ -315,7 +337,7 @@ def list_directory(dirname): dirname: string, path to a directory Returns: - [filename1, filename2, ... filenameN] + [filename1, filename2, ... filenameN] as strings Raises: errors.NotFoundError if directory doesn't exist @@ -323,8 +345,10 @@ def list_directory(dirname): if not is_directory(dirname): raise errors.NotFoundError(None, None, "Could not find directory") file_list = get_matching_files(os.path.join(compat.as_str_any(dirname), "*")) - return [compat.as_bytes(pywrap_tensorflow.Basename(compat.as_bytes(filename))) - for filename in file_list] + return [ + compat.as_str_any(pywrap_tensorflow.Basename(compat.as_bytes(filename))) + for filename in file_list + ] def walk(top, in_order=True): @@ -337,11 +361,12 @@ def walk(top, in_order=True): Errors that happen while listing directories are ignored. Yields: - # Each yield is a 3-tuple: the pathname of a directory, followed - # by lists of all its subdirectories and leaf files. + Each yield is a 3-tuple: the pathname of a directory, followed by lists of + all its subdirectories and leaf files. (dirname, [subdirname, subdirname, ...], [filename, filename, ...]) + as strings """ - top = compat.as_bytes(top) + top = compat.as_str_any(top) try: listing = list_directory(top) except errors.NotFoundError: |