⚙️ update graph classes

src/defigraph/Edge.py

@@ -1,15 +1,17 @@
 from defigraph.Vertex import Vertex
 from defigraph.Pool import Pool
+from math import log
 
 class Edge:
-  def __init__(self, u: Vertex, v: Vertex, pool_address: str, weight: float=None):
-    self.u = u
-    self.v = v
-    self.weight = weight
-    self.pool = Pool(address=pool_address,token0=u, token1=v)
+  def __init__(self, pool: Pool):
+    self.pool = pool
+    self.weight = -log(self.pool.token0_price)
 
   def __repr__(self):
-    if not self.weight:
-      return f"({self.u}, {self.v})"
-    else:
-      return f"({self.u}, {self.v}, {self.weight})"
+    return f"({self.pool.token0}, {self.pool.token1}, {self.weight})"
+
+  def __eq__(self, edge):
+    return self.pool == edge.pool
+
+  def __ne__(self, edge):
+    return self.pool != edge.pool

src/defigraph/Graph.py

@@ -1,3 +1,89 @@
+from defigraph.Vertex import Vertex
+from defigraph.Edge import Edge
+from defigraph.Pool import Pool
+from typing import List, Dict
+
 class Graph:
-  def __init__(self):
-    pass
+  def __init__(self, edges: List[Edge]):
+    self.vertices: List[Vertex] = []
+    self.adjascency_list: Dict[Edge] = {}
+
+    for edge in edges:
+      # Create adjacency list
+      if edge.pool.token0.name not in self.adjascency_list:
+        self.adjascency_list[edge.pool.token0.name] = [edge]
+      else:
+        if edge not in self.adjascency_list[edge.pool.token0.name]:
+          self.adjascency_list[edge.pool.token0.name].append(edge)
+      # Create vertices list
+      if edge.pool.token0 not in self.vertices:
+        self.vertices.append(edge.pool.token0)
+      if edge.pool.token1 not in self.vertices:
+        self.vertices.append(edge.pool.token1)
+
+  def __repr__(self):
+    return f"{self.adjascency_list}"
+
+  def __getitem__(self, vertex: str):
+    return self.adjascency_list[vertex]
+
+  def __setitem__(self, vertex: str, edges: List[Vertex]):
+    if vertex not in self.adjascency_list:
+      self.adjascency_list[vertex] = edges
+    else:
+      for edge in edges:
+        if edge not in self.adjascency_list[vertex]:
+          self.adjascency_list[vertex].append(edge)
+
+  def __len__(self):
+    return len(self.vertices)
+
+  def add_edge(self, pool: Pool, directed=False):
+    """Directed=False means you can trade in any direction between token0 and token1"""
+    edge = Edge(pool=pool)
+
+    if edge.pool.token0.name not in self.adjascency_list:
+      self.adjascency_list[edge.pool.token0.name] = [edge]
+    else:
+      if edge not in self.adjascency_list[edge.pool.token0.name]:
+        self.adjascency_list[edge.pool.token0.name].append(edge)
+
+    # add the undirected edge
+    if not directed:
+      token0 = edge.pool.token0
+      token1 = edge.pool.token1
+      pool_address = edge.pool.address
+      fee = edge.pool.fee
+      token0price = edge.pool.token0_price
+      token1price = edge.pool.token1_price
+      reverse_pool_direction = Pool(token0=token1,token1=token0, pool_address=pool_address, fee=fee, token0_price=token1price, token1_price=token0price)
+      edge = Edge(pool=reverse_pool_direction)
+
+      if edge.pool.token0.name not in self.adjascency_list:
+        self.adjascency_list[edge.pool.token0.name] = [edge]
+      else:
+        if edge not in self.adjascency_list[edge.pool.token0.name]:
+          self.adjascency_list[edge.pool.token0.name].append(edge)
+
+    # Update vertices list
+    if edge.pool.token0 not in self.vertices:
+      self.vertices.append(edge.pool.token0)
+    if edge.pool.token1 not in self.vertices:
+      self.vertices.append(edge.pool.token1)
+
+  def get_edge_count(self):
+    count = 0
+    for vertex in self.adjascency_list:
+      count += len(self.adjascency_list[vertex])
+    return count
+
+  def get_vertices(self):
+    return self.vertices
+
+  def get_edges(self):
+    edges = []
+    for vertex in self.adjascency_list:
+      edges.append(self.adjascency_list[vertex])
+    return edges
+
+

src/defigraph/Pool.py

@@ -1,7 +1,22 @@
 from defigraph.Vertex import Vertex
 
 class Pool:
-  def __init__(self, address: str, token0: Vertex, token1: Vertex):
-    self.address = address
+  def __init__(self, pool_address: str, token0: Vertex, token1: Vertex, fee: int, token0_price: float, token1_price: float):
+    self.address = pool_address
     self.token0 = token0
-    self.token1 = token1
+    self.token1 = token1
+    self.token0_price = token0_price
+    self.token1_price = token1_price
+    self.fee = fee
+
+  def __repr__(self):
+    return f"{(self.token0, self.token1, self.fee)}"
+
+  def __eq__(self, pool):
+    return self.address == pool.address and self.fee == pool.fee and (self.token0 == pool.token0 or self.token0 == pool.token1) and (self.token1 == pool.token1 or self.token1 == pool.token0)
+
+  def __ne__(self, pool):
+    return self.address != pool.address or self.fee != pool.fee
+
+  def __hash__(self):
+    return hash(str(self))

src/defigraph/Vertex.py

@@ -1,9 +1,17 @@
 class Vertex:
-  def __init__(self, name: str, decimals: int):
+  def __init__(self, name: str, decimals: int, address: str):
     self.name = name
     self.decimals = decimals
+    self.address = address
 
   def __repr__(self):
     return f"{self.name}"
-
-
+
+  def __eq__(self, vertex):
+    return self.name == vertex.name and self.decimals == vertex.decimals and self.address == vertex.address
+
+  def __ne__(self, vertex):
+    return self.name != vertex.name or self.decimals != vertex.decimals or self.address != vertex.address
+
+  def __hash__(self):
+    return hash(str(self))

src/defigraph/Edge_test.py → tests/Edge_test.py

@@ -1,7 +1,7 @@
 import unittest
 
 class TestEdgeMethods(unittest.TestCase):
-  """Test class for Edge methods"""
+  pass
 
 if __name__ == "__main__":
   unittest.main()

src/defigraph/Graph_test.py → tests/Graph_test.py

@@ -1,4 +1,7 @@
 import unittest
 
 class TestGraphMethods(unittest.TestCase):
-  """Test class for Graph methods"""
+  pass
+
+if __name__ == "__main__":
+  unittest.main()

tests/Pool_test.py

@@ -0,0 +1,7 @@
+import unittest
+
+class TestPoolMethods(unittest.TestCase):
+  pass
+
+if __name__ == "__main__":
+  unittest.main()

src/defigraph/Vertex_test.py → tests/Vertex_test.py

@@ -2,19 +2,25 @@
 from defigraph.Vertex import Vertex
 
 class TestVertexMethods(unittest.TestCase):
-  """Test class for Edge methods"""
-
   def setUp(self):
     name = "WETH"
     decimals = 18
-    self.vertex = Vertex(name=name,decimals=decimals)
+    addresss = "0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2"
+    self.vertex = Vertex(name=name,decimals=decimals, address=addresss)
 
   def test_vertex_has_name(self):
+
     assert self.vertex.name is not None
 
   def test_vertex_has_decimals(self):
     assert self.vertex.decimals is not None
 
+  def test_vertex_has_address(self):
+    assert self.vertex.address is not None
+
+  def test_address_length(self):
+    assert len(self.vertex.address) == 42
+
   def tearDown(self):
     return super().tearDown()