From dd394f88617505458f34d929c3c472dca65dad00 Mon Sep 17 00:00:00 2001
From: Darcy LF <ehnevault@gmail.com>
Date: Sat, 14 Aug 2021 10:45:35 +1000
Subject: [PATCH] :sparkles: add edge.set_weight() and edge.weight()

---
 pynode_next/edge.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/pynode_next/edge.py b/pynode_next/edge.py
index d12c440..d8c4b1a 100644
--- a/pynode_next/edge.py
+++ b/pynode_next/edge.py
@@ -12,7 +12,7 @@ def __init__(self, source, target, weight=None, directed=False):
         self._weight = weight
         if weight == None:
             self._weight = ""
-        
+
         self._color = Color.LIGHT_GREY
 
         self._internal_id = uuid.uuid4()
@@ -70,6 +70,18 @@ def color(self):
         """Gets the edge's color."""
         return self._color
 
+    def set_weight(self, weight):
+        """Sets the edge's weight. (Weight must be serialisable)"""
+        self._weight = weight
+        core.ax(
+            lambda x: self._dispatch_wrapper(x, {"labels": {1: {"text": str(weight)}}})
+        )
+        return self
+
+    def weight(self):
+        """Returns the edge's weight. Returns an empty string if weight wasn't defined at init and hasn't been changed since."""
+        return self._weight
+
     def traverse(self, initial_node=None, color=Color.RED, keep_path=True):
         if initial_node == None:
             source = self._source
@@ -114,11 +126,7 @@ def _data(self):
                         "source": str(self._source),
                         "target": str(self._target),
                         "directed": self._directed,
-                        "labels": {
-                            1: {
-                                "text": str(self._weight)
-                            }
-                        }
+                        "labels": {1: {"text": str(self._weight)}},
                     }
                 }
             }