From 71d4679caf4564398a389c56718d0510909d37ab Mon Sep 17 00:00:00 2001
From: jeffthibault <jdthibault2@gmail.com>
Date: Wed, 3 Aug 2022 12:10:02 -0400
Subject: [PATCH] refactor relay_manager and add remove_relay method

---
 nostr/relay_manager.py | 49 ++++++++++++++++++++++++++++++------------
 1 file changed, 35 insertions(+), 14 deletions(-)

diff --git a/nostr/relay_manager.py b/nostr/relay_manager.py
index 897f577..0dc0c01 100644
--- a/nostr/relay_manager.py
+++ b/nostr/relay_manager.py
@@ -5,35 +5,56 @@ from .relay import Relay, RelayPolicy
 
 class RelayManager:
     def __init__(self) -> None:
-        self.relays: list[Relay] = []
+        self.relays: dict[str, Relay] = {}
         self.message_pool = MessagePool()
 
     def add_relay(self, url: str, read: bool, write: bool, subscriptions={}):
         policy = RelayPolicy(read, write)
         relay = Relay(url, policy, self.message_pool, subscriptions)
-        self.relays.append(relay)
+        self.relays[url] = relay
 
-    def add_subscription(self, id: str, filters: Filters):
-        for relay in self.relays:
+    def remove_relay(self, url: str):
+        self.relays.pop(url)
+
+    def add_subscription(self, id: str, filters: Filters, relay: Relay=None):
+        if relay != None:
             relay.add_subscription(id, filters)
+        else:
+            for relay in self.relays.values():
+                relay.add_subscription(id, filters)
 
-    def close_subscription(self, id: str):
-        for relay in self.relays:
+    def close_subscription(self, id: str, relay: Relay=None):
+        if relay != None:
             relay.close_subscription(id)
+        else:
+            for relay in self.relays.values():
+                relay.close_subscription(id)
 
-    def open_connections(self):
-        for relay in self.relays:
+    def open_connection(self, relay: Relay=None):
+        if relay != None:
             threading.Thread(
                 target=relay.connect,
                 name=f"{relay.url}-thread"
             ).start()
+        else:
+            for relay in self.relays.values():
+                threading.Thread(
+                    target=relay.connect,
+                    name=f"{relay.url}-thread"
+                ).start()
 
-    def close_connections(self):
-        for relay in self.relays:
+    def close_connection(self, relay: Relay=None):
+        if relay != None:
             relay.close()
+        else:
+            for relay in self.relays.values():
+                relay.close()
 
-    def publish_message(self, message: str):
-        for relay in self.relays:
-            if relay.policy.should_write:
-                relay.publish(message)
+    def publish_message(self, message: str, relay: Relay=None):
+        if relay != None:
+            relay.publish(message)
+        else:
+            for relay in self.relays.values():
+                if relay.policy.should_write:
+                    relay.publish(message)