#! /usr/bin/env python3

import os

import tensorflow as tf


def wr(g, name):
    name = os.path.join(os.path.dirname(__file__), "{}.pb".format(name))
    open(name, "wb").write(g.as_graph_def().SerializeToString())


g = tf.Graph()
with g.as_default():
    c = tf.constant(123.0)
wr(g, "constant")

@@ -27,16 +27,54 @@
 mport os
+%0Afrom contextlib import contextmanager
 %0A%0Aimport
@@ -97,18 +97,32 @@
 f%0A%0A%0A
-def wr(g, 
+@contextmanager%0Adef gen(
 name
@@ -201,16 +201,21 @@
 e))%0A    
+with 
 open(nam
@@ -222,16 +222,106 @@
 e, %22wb%22)
+ as out:%0A        g = tf.Graph()%0A        with g.as_default():%0A            yield%0A        out
 .write(g
@@ -363,51 +363,33 @@
 )%0A%0A%0A
-g = tf.Graph()%0Awith g.as_default(
+with gen(%22constant%22
 ):%0A   
- c =
  tf.
@@ -408,22 +408,4 @@
 .0)%0A
-wr(g, %22constant%22)%0A

Refactor gengraph.py