Adds Operations() method to Graph

There is currently no way to list all of the operations in a graph
from the go api. This patch ads an Operations() method to retrieve the
list using the existing TF_GraphNextOperation c api. The graph_test
was modified to include testing this new method.

Signed-off-by: Vishvananda Ishaya Abrams <vishvananda@gmail.com>
This commit is contained in:
Vishvananda Ishaya Abrams 2017-11-29 22:25:37 -08:00
parent 8a4d849691
commit 5d52b95279
2 changed files with 33 additions and 3 deletions

View File

@ -114,6 +114,20 @@ func (g *Graph) Operation(name string) *Operation {
return &Operation{cop, g}
}
// Operations returns a list of all operations in the graph
func (g *Graph) Operations() []Operation {
var pos C.size_t = 0
ops := []Operation{}
for {
cop := C.TF_GraphNextOperation(g.c, &pos)
if cop == nil {
break
}
ops = append(ops, Operation{cop, g})
}
return ops
}
// OpSpec is the specification of an Operation to be added to a Graph
// (using Graph.AddOperation).
type OpSpec struct {

View File

@ -29,10 +29,26 @@ func hasOperations(g *Graph, ops ...string) error {
missing = append(missing, op)
}
}
if len(missing) == 0 {
return nil
if len(missing) != 0 {
return fmt.Errorf("Graph does not have the operations %v", missing)
}
return fmt.Errorf("Graph does not have the operations %v", missing)
inList := map[string]bool{}
for _, op := range g.Operations() {
inList[op.Name()] = true
}
for _, op := range ops {
if !inList[op] {
missing = append(missing, op)
}
}
if len(missing) != 0 {
return fmt.Errorf("Operations %v are missing from graph.Operations()", missing)
}
return nil
}
func TestGraphWriteToAndImport(t *testing.T) {