From fd26e49734260f72290c0cb391a31444519fe4e1 Mon Sep 17 00:00:00 2001
From: filifa <filifa.breath652@8alias.com>
Date: Tue, 6 May 2025 23:56:20 -0400
Subject: [PATCH] add argument for setting weight attribute

---
 cmd/internal/graph/dot/graph.go       | 7 ++++---
 cmd/internal/graph/dot/nodes_edges.go | 3 ++-
 cmd/root.go                           | 6 +++++-
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/cmd/internal/graph/dot/graph.go b/cmd/internal/graph/dot/graph.go
index cff53d8..ecdc50b 100644
--- a/cmd/internal/graph/dot/graph.go
+++ b/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.
diff --git a/cmd/internal/graph/dot/nodes_edges.go b/cmd/internal/graph/dot/nodes_edges.go
index f13a030..d876346 100644
--- a/cmd/internal/graph/dot/nodes_edges.go
+++ b/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)
diff --git a/cmd/root.go b/cmd/root.go
index 942aabd..cd9a692 100644
--- a/cmd/root.go
+++ b/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")
 }