Have write_graph return the output path of file (tensorflow#6851)

samjabrahams · drpngx · commit 48c433f48bb3 · 2017-01-14T12:18:58.000-08:00
* Have `write_graph` return the output path of file

Since the function is performing `os.path.join()` for the user, it makes
sense to return the location of the final file, or else the user will
have to call `os.path.join()` again on their own.

* Adjust WriteGraphTest to check returned path

* Assert path exists; "/".join -&gt; os.path.join

Theoretically more Windows compatible
diff --git a/tensorflow/python/framework/graph_io.py b/tensorflow/python/framework/graph_io.py
@@ -50,6 +50,9 @@ def write_graph(graph_or_graph_def, logdir, name, as_text=True):
       filesystems, such as Google Cloud Storage (GCS).
     name: Filename for the graph.
     as_text: If `True`, writes the graph as an ASCII proto.
+
+  Returns:
+    The path of the output proto file.
   """
   if isinstance(graph_or_graph_def, ops.Graph):
     graph_def = graph_or_graph_def.as_graph_def()
@@ -64,3 +67,4 @@ def write_graph(graph_or_graph_def, logdir, name, as_text=True):
     file_io.atomic_write_string_to_file(path, str(graph_def))
   else:
     file_io.atomic_write_string_to_file(path, graph_def.SerializeToString())
+  return path
diff --git a/tensorflow/python/training/saver_test.py b/tensorflow/python/training/saver_test.py
@@ -1837,14 +1837,22 @@ class WriteGraphTest(test.TestCase):
   def testWriteGraph(self):
     test_dir = _TestDir("write_graph_dir")
     variables.Variable([[1, 2, 3], [4, 5, 6]], dtype=dtypes.float32, name="v0")
-    graph_io.write_graph(ops_lib.get_default_graph(),
-                         "/".join([test_dir, "l1"]), "graph.pbtxt")
+    path = graph_io.write_graph(ops_lib.get_default_graph(),
+                                os.path.join(test_dir, "l1"), "graph.pbtxt")
+    truth = os.path.join(test_dir, "l1", "graph.pbtxt")
+    self.assertEqual(path, truth)
+    self.assertTrue(os.path.exists(path))
+
 
   def testRecursiveCreate(self):
     test_dir = _TestDir("deep_dir")
     variables.Variable([[1, 2, 3], [4, 5, 6]], dtype=dtypes.float32, name="v0")
-    graph_io.write_graph(ops_lib.get_default_graph().as_graph_def(),
-                         "/".join([test_dir, "l1/l2/l3"]), "graph.pbtxt")
+    path = graph_io.write_graph(ops_lib.get_default_graph().as_graph_def(),
+                                os.path.join(test_dir, "l1", "l2", "l3"),
+                                "graph.pbtxt")
+    truth = os.path.join(test_dir, 'l1', 'l2', 'l3', "graph.pbtxt")
+    self.assertEqual(path, truth)
+    self.assertTrue(os.path.exists(path))
 
 
 class SaverUtilsTest(test.TestCase):
