Adding RedisGraph support (#1673)

Co-authored-by: Chayim I. Kirshen <c@kirshen.com>

1 files changed, 87 insertions, 0 deletions

diff --git a/redis/commands/graph/edge.py b/redis/commands/graph/edge.py
new file mode 100644
index 0000000..b334293
--- /dev/null
+++ b/redis/commands/graph/edge.py
@@ -0,0 +1,87 @@
+from ..helpers import quote_string
+from .node import Node
+
+
+class Edge:
+    """
+    An edge connecting two nodes.
+    """
+
+    def __init__(self, src_node, relation, dest_node, edge_id=None, properties=None):
+        """
+        Create a new edge.
+        """
+        if src_node is None or dest_node is None:
+            # NOTE(bors-42): It makes sense to change AssertionError to
+            #                ValueError here
+            raise AssertionError("Both src_node & dest_node must be provided")
+
+        self.id = edge_id
+        self.relation = relation or ""
+        self.properties = properties or {}
+        self.src_node = src_node
+        self.dest_node = dest_node
+
+    def toString(self):
+        res = ""
+        if self.properties:
+            props = ",".join(
+                key + ":" + str(quote_string(val))
+                for key, val in sorted(self.properties.items())
+            )
+            res += "{" + props + "}"
+
+        return res
+
+    def __str__(self):
+        # Source node.
+        if isinstance(self.src_node, Node):
+            res = str(self.src_node)
+        else:
+            res = "()"
+
+        # Edge
+        res += "-["
+        if self.relation:
+            res += ":" + self.relation
+        if self.properties:
+            props = ",".join(
+                key + ":" + str(quote_string(val))
+                for key, val in sorted(self.properties.items())
+            )
+            res += "{" + props + "}"
+        res += "]->"
+
+        # Dest node.
+        if isinstance(self.dest_node, Node):
+            res += str(self.dest_node)
+        else:
+            res += "()"
+
+        return res
+
+    def __eq__(self, rhs):
+        # Quick positive check, if both IDs are set.
+        if self.id is not None and rhs.id is not None and self.id == rhs.id:
+            return True
+
+        # Source and destination nodes should match.
+        if self.src_node != rhs.src_node:
+            return False
+
+        if self.dest_node != rhs.dest_node:
+            return False
+
+        # Relation should match.
+        if self.relation != rhs.relation:
+            return False
+
+        # Quick check for number of properties.
+        if len(self.properties) != len(rhs.properties):
+            return False
+
+        # Compare properties.
+        if self.properties != rhs.properties:
+            return False
+
+        return True