add argument for setting weight attribute

cmd/internal/graph/dot/graph.go

@@ -28,17 +28,18 @@ import (
 // DOTDirectedGraph is a graph to unmarshal DOT graphs.
 type DOTDirectedGraph struct {
 	*igraph.DirectedGraph
+	WeightAttribute string
 }
 
 // NewDOTDirectedGraph returns a graph with no nodes or edges.
-func NewDOTDirectedGraph() *DOTDirectedGraph {
-	return &DOTDirectedGraph{DirectedGraph: igraph.NewDirectedGraph()}
+func NewDOTDirectedGraph(weightAttr string) *DOTDirectedGraph {
+	return &DOTDirectedGraph{DirectedGraph: igraph.NewDirectedGraph(), WeightAttribute: weightAttr}
 }
 
 // NewLine returns a DOT-aware weighted line.
 func (g *DOTDirectedGraph) NewLine(from, to graph.Node) graph.Line {
 	e := g.DirectedGraph.NewWeightedLine(from, to, math.NaN()).(multi.WeightedLine)
-	return &weightedLine{WeightedLine: e}
+	return &weightedLine{WeightedLine: e, weightAttribute: g.WeightAttribute}
 }
 
 // NewNode returns a DOT-aware Node.

cmd/internal/graph/dot/nodes_edges.go

@@ -46,6 +46,7 @@ func (n *Node) DOTID() string {
 // multi.WeightedLine, it allows for self-loops.
 type weightedLine struct {
 	multi.WeightedLine
+	weightAttribute string
 }
 
 // SetAttribute enables storing the weight read from a DOT file. It errors if
@@ -54,7 +55,7 @@ func (e *weightedLine) SetAttribute(attr encoding.Attribute) error {
 	var err error
 
 	switch attr.Key {
-	case "len":
+	case e.weightAttribute:
 		e.W, err = strconv.ParseFloat(attr.Value, 64)
 	default:
 		err = errors.New("unknown key:" + attr.Key)

cmd/root.go

@@ -35,6 +35,8 @@ var pythonFmt bool
 
 var oneline bool
 
+var weightAttr string
+
 var nodeOrder []string
 
 // rootCmd represents the base command when called without any subcommands
@@ -58,7 +60,7 @@ func parse(cmd *cobra.Command, args []string) {
 		panic(err)
 	}
 
-	graph := idot.NewDOTDirectedGraph()
+	graph := idot.NewDOTDirectedGraph(weightAttr)
 	err = dot.UnmarshalMulti(data, graph)
 	if err != nil {
 		panic(err)
@@ -148,5 +150,7 @@ func init() {
 
 	rootCmd.Flags().BoolVar(&oneline, "oneline", false, "output on one line")
 
+	rootCmd.Flags().StringVar(&weightAttr, "weight-attr", "len", "edge attribute to use as weight")
+
 	rootCmd.Flags().StringSliceVarP(&nodeOrder, "order", "o", nil, "order of nodes in rows/columns of output")
 }