Update asset manager filesystem to handle directories correctly.

Change: 141478190

2 files changed, 37 insertions, 4 deletions

diff --git a/tensorflow/contrib/android/asset_manager_filesystem.cc b/tensorflow/contrib/android/asset_manager_filesystem.cc
index 49d179605c..9e4d3290c3 100644
--- a/tensorflow/contrib/android/asset_manager_filesystem.cc
+++ b/tensorflow/contrib/android/asset_manager_filesystem.cc
@@ -17,11 +17,19 @@ limitations under the License.
 
 #include <unistd.h>
 
+#include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/env.h"
 
 namespace tensorflow {
 namespace {
 
+string RemoveSuffix(const string& name, const string& suffix) {
+  string output(name);
+  StringPiece piece(output);
+  str_util::ConsumeSuffix(&piece, suffix);
+  return piece.ToString();
+}
+
 // Closes the given AAsset when variable is destructed.
 class ScopedAsset {
  public:
@@ -174,23 +182,29 @@ Status AssetManagerFileSystem::NewReadOnlyMemoryRegionFromFile(
   return Status::OK();
 }
 
-Status AssetManagerFileSystem::GetChildren(const string& dir_name,
+Status AssetManagerFileSystem::GetChildren(const string& prefixed_dir,
                                            std::vector<string>* r) {
-  string path = RemoveAssetPrefix(dir_name);
+  std::string path = NormalizeDirectoryPath(prefixed_dir);
   auto dir =
       ScopedAssetDir(AAssetManager_openDir(asset_manager_, path.c_str()));
   if (dir.get() == nullptr) {
-    return errors::NotFound("Directory ", dir_name, " not found.");
+    return errors::NotFound("Directory ", prefixed_dir, " not found.");
   }
   const char* next_file = AAssetDir_getNextFileName(dir.get());
   while (next_file != nullptr) {
-    r->push_back(string(next_file));
+    r->push_back(next_file);
     next_file = AAssetDir_getNextFileName(dir.get());
   }
   return Status::OK();
 }
 
 Status AssetManagerFileSystem::GetFileSize(const string& fname, uint64* s) {
+  // If fname corresponds to a directory, return early. It doesn't map to an
+  // AAsset, and would otherwise return NotFound.
+  if (DirectoryExists(fname)) {
+    *s = 0;
+    return Status::OK();
+  }
   string path = RemoveAssetPrefix(fname);
   auto asset = ScopedAsset(
       AAssetManager_open(asset_manager_, path.c_str(), AASSET_MODE_RANDOM));
@@ -203,11 +217,16 @@ Status AssetManagerFileSystem::GetFileSize(const string& fname, uint64* s) {
 
 Status AssetManagerFileSystem::Stat(const string& fname, FileStatistics* stat) {
   uint64 size;
+  stat->is_directory = DirectoryExists(fname);
   TF_RETURN_IF_ERROR(GetFileSize(fname, &size));
   stat->length = size;
   return Status::OK();
 }
 
+string AssetManagerFileSystem::NormalizeDirectoryPath(const string& fname) {
+  return RemoveSuffix(RemoveAssetPrefix(fname), "/");
+}
+
 string AssetManagerFileSystem::RemoveAssetPrefix(const string& name) {
   string output(name);
   StringPiece piece(output);
@@ -215,6 +234,15 @@ string AssetManagerFileSystem::RemoveAssetPrefix(const string& name) {
   return piece.ToString();
 }
 
+bool AssetManagerFileSystem::DirectoryExists(const std::string& fname) {
+  std::string path = NormalizeDirectoryPath(fname);
+  auto dir =
+      ScopedAssetDir(AAssetManager_openDir(asset_manager_, path.c_str()));
+  // Note that openDir will return something even if the directory doesn't
+  // exist. Therefore, we need to ensure one file exists in the folder.
+  return AAssetDir_getNextFileName(dir.get()) != NULL;
+}
+
 Status AssetManagerFileSystem::NewWritableFile(
     const string& fname, std::unique_ptr<WritableFile>* result) {
   return errors::Unimplemented("Asset storage is read only.");
diff --git a/tensorflow/contrib/android/asset_manager_filesystem.h b/tensorflow/contrib/android/asset_manager_filesystem.h
index 8cc31db9f2..2b43939f14 100644
--- a/tensorflow/contrib/android/asset_manager_filesystem.h
+++ b/tensorflow/contrib/android/asset_manager_filesystem.h
@@ -69,6 +69,11 @@ class AssetManagerFileSystem : public FileSystem {
  private:
   string RemoveAssetPrefix(const string& name);
 
+  // Return a string path that can be passed into AAssetManager functions.
+  // For example, 'my_prefix://some/dir/' would return 'some/dir'.
+  string NormalizeDirectoryPath(const string& fname);
+  bool DirectoryExists(const std::string& fname);
+
   AAssetManager* asset_manager_;
   string prefix_;
 };