Merge pull request #14 from BlockScience/refactor

v1.1.0-beta.3

Author: lukvmil

examples/partial.py

@@ -17,7 +17,7 @@
 
 
 class PartialNodeConfig(NodeConfig):
-    koi_net: KoiNetConfig | None = Field(default_factory = lambda:
+    koi_net: KoiNetConfig = Field(default_factory = lambda:
         KoiNetConfig(
             node_name="partial",
             node_profile=NodeProfile(
@@ -34,4 +34,4 @@ class PartialNodeConfig(NodeConfig):
 )
 
 if __name__ == "__main__":
-    node.poller.run()
+    node.lifecycle.start()

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "koi-net"
-version = "1.1.0-beta.2"
+version = "1.1.0-beta.3"
 description = "Implementation of KOI-net protocol in Python"
 authors = [
     {name = "Luke Miller", email = "[email protected]"}

src/koi_net/core.py

@@ -1,5 +1,5 @@
 import logging
-from koi_net.protocol.node import NodeType
+from typing import Generic, TypeVar
 from rid_lib.ext import Cache
 from .network.resolver import NetworkResolver
 from .network.event_queue import NetworkEventQueue
@@ -25,9 +25,10 @@
 logger = logging.getLogger(__name__)
 
 
+T = TypeVar("T", bound=NodeConfig)
 
-class NodeInterface:
-    config: NodeConfig
+class NodeInterface(Generic[T]):
+    config: T
     cache: Cache
     identity: NodeIdentity
     resolver: NetworkResolver
@@ -41,58 +42,75 @@ class NodeInterface:
 
     def __init__(
         self,
-        config: NodeConfig,
+        config: T,
         use_kobj_processor_thread: bool = False,
         handlers: list[KnowledgeHandler] | None = None,
-        cache: Cache | None = None,
-        processor: ProcessorInterface | None = None
+        
+        CacheOverride: type[Cache] | None = None,
+        NodeIdentityOverride: type[NodeIdentity] | None = None,
+        EffectorOverride: type[Effector] | None = None,
+        NetworkGraphOverride: type[NetworkGraph] | None = None,
+        SecureOverride: type[Secure] | None = None,
+        RequestHandlerOverride: type[RequestHandler] | None = None,
+        ResponseHandlerOverride: type[ResponseHandler] | None = None,
+        NetworkResolverOverride: type[NetworkResolver] | None = None,
+        NetworkEventQueueOverride: type[NetworkEventQueue] | None = None,
+        ActorOverride: type[Actor] | None = None,
+        ActionContextOverride: type[ActionContext] | None = None,
+        HandlerContextOverride: type[HandlerContext] | None = None,
+        KnowledgePipelineOverride: type[KnowledgePipeline] | None = None,
+        ProcessorInterfaceOverride: type[ProcessorInterface] | None = None,
+        ErrorHandlerOverride: type[ErrorHandler] | None = None,
+        NodeLifecycleOverride: type[NodeLifecycle] | None = None,
+        NodeServerOverride: type[NodeServer] | None = None,
+        NodePollerOverride: type[NodePoller] | None = None,        
     ):
         self.config = config
-        self.cache = cache or Cache(
+        self.cache = (CacheOverride or Cache)(
             directory_path=self.config.koi_net.cache_directory_path
         )
 
-        self.identity = NodeIdentity(config=self.config)
-        self.effector = Effector(cache=self.cache)
+        self.identity = (NodeIdentityOverride or NodeIdentity)(config=self.config)
+        self.effector = (EffectorOverride or Effector)(cache=self.cache)
 
-        self.graph = NetworkGraph(
-            cache=self.cache, 
+        self.graph = (NetworkGraphOverride or NetworkGraph)(
+            cache=self.cache,
             identity=self.identity
         )
-        
-        self.secure = Secure(
-            identity=self.identity, 
-            effector=self.effector, 
+
+        self.secure = (SecureOverride or Secure)(
+            identity=self.identity,
+            effector=self.effector,
             config=self.config
         )
-        
-        self.request_handler = RequestHandler(
-            effector=self.effector, 
+
+        self.request_handler = (RequestHandlerOverride or RequestHandler)(
+            effector=self.effector,
             identity=self.identity,
             secure=self.secure
         )
-        
-        self.response_handler = ResponseHandler(self.cache, self.effector)
-        
-        self.resolver = NetworkResolver(
+
+        self.response_handler = (ResponseHandlerOverride or ResponseHandler)(self.cache, self.effector)
+
+        self.resolver = (NetworkResolverOverride or NetworkResolver)(
             config=self.config,
-            cache=self.cache, 
+            cache=self.cache,
             identity=self.identity,
             graph=self.graph,
             request_handler=self.request_handler,
             effector=self.effector
         )
-        
-        self.event_queue = NetworkEventQueue(
+
+        self.event_queue = (NetworkEventQueueOverride or NetworkEventQueue)(
             config=self.config,
-            cache=self.cache, 
+            cache=self.cache,
             identity=self.identity,
             graph=self.graph,
             request_handler=self.request_handler,
             effector=self.effector
         )
 
-        self.actor = Actor(
+        self.actor = (ActorOverride or Actor)(
             identity=self.identity,
             effector=self.effector,
             event_queue=self.event_queue
@@ -107,12 +125,12 @@ def __init__(
 
         self.use_kobj_processor_thread = use_kobj_processor_thread
 
-        self.action_context = ActionContext(
+        self.action_context = (ActionContextOverride or ActionContext)(
             identity=self.identity,
             effector=self.effector
         )
 
-        self.handler_context = HandlerContext(
+        self.handler_context = (HandlerContextOverride or HandlerContext)(
             identity=self.identity,
             cache=self.cache,
             event_queue=self.event_queue,
@@ -121,7 +139,7 @@ def __init__(
             effector=self.effector
         )
 
-        self.pipeline = KnowledgePipeline(
+        self.pipeline = (KnowledgePipelineOverride or KnowledgePipeline)(
             handler_context=self.handler_context,
             cache=self.cache,
             request_handler=self.request_handler,
@@ -130,12 +148,12 @@ def __init__(
             default_handlers=handlers
         )
 
-        self.processor = processor or ProcessorInterface(
+        self.processor = (ProcessorInterfaceOverride or ProcessorInterface)(
             pipeline=self.pipeline,
             use_kobj_processor_thread=self.use_kobj_processor_thread
         )
 
-        self.error_handler = ErrorHandler(
+        self.error_handler = (ErrorHandlerOverride or ErrorHandler)(
             processor=self.processor,
             actor=self.actor
         )
@@ -148,7 +166,7 @@ def __init__(
         self.effector.set_resolver(self.resolver)
         self.effector.set_action_context(self.action_context)
 
-        self.lifecycle = NodeLifecycle(
+        self.lifecycle = (NodeLifecycleOverride or NodeLifecycle)(
             config=self.config,
             identity=self.identity,
             graph=self.graph,
@@ -159,7 +177,7 @@ def __init__(
         )
 
         # if self.config.koi_net.node_profile.node_type == NodeType.FULL:
-        self.server = NodeServer(
+        self.server = (NodeServerOverride or NodeServer)(
             config=self.config,
             lifecycle=self.lifecycle,
             secure=self.secure,
@@ -168,7 +186,7 @@ def __init__(
             response_handler=self.response_handler
         )
 
-        self.poller = NodePoller(
+        self.poller = (NodePollerOverride or NodePoller)(
             processor=self.processor,
             lifecycle=self.lifecycle,
             resolver=self.resolver,

src/koi_net/network/response_handler.py

@@ -1,5 +1,6 @@
 import logging
 from rid_lib import RID
+from rid_lib.types import KoiNetNode
 from rid_lib.ext import Manifest, Cache
 from rid_lib.ext.bundle import Bundle
 
@@ -30,13 +31,13 @@ def __init__(
         self.cache = cache
         self.effector = effector
 
-    def fetch_rids(self, req: FetchRids) -> RidsPayload:
+    def fetch_rids(self, req: FetchRids, source: KoiNetNode) -> RidsPayload:
         logger.info(f"Request to fetch rids, allowed types {req.rid_types}")
         rids = self.cache.list_rids(req.rid_types)
 
         return RidsPayload(rids=rids)
 
-    def fetch_manifests(self, req: FetchManifests) -> ManifestsPayload:
+    def fetch_manifests(self, req: FetchManifests, source: KoiNetNode) -> ManifestsPayload:
         logger.info(f"Request to fetch manifests, allowed types {req.rid_types}, rids {req.rids}")
 
         manifests: list[Manifest] = []
@@ -51,7 +52,7 @@ def fetch_manifests(self, req: FetchManifests) -> ManifestsPayload:
 
         return ManifestsPayload(manifests=manifests, not_found=not_found)
 
-    def fetch_bundles(self, req: FetchBundles) -> BundlesPayload:
+    def fetch_bundles(self, req: FetchBundles, source: KoiNetNode) -> BundlesPayload:
         logger.info(f"Request to fetch bundles, requested rids {req.rids}")
 
         bundles: list[Bundle] = []

src/koi_net/poller.py

@@ -22,17 +22,19 @@ def __init__(
         self.resolver = resolver
         self.config = config
 
+    def poll(self):
+        neighbors = self.resolver.poll_neighbors()
+        for node_rid in neighbors:
+            for event in neighbors[node_rid]:
+                self.processor.handle(event=event, source=node_rid)
+        self.processor.flush_kobj_queue()
+
     def run(self):
         try:
             self.lifecycle.start()
             while True:
                 start_time = time.time()
-                neighbors = self.resolver.poll_neighbors()
-                for node_rid in neighbors:
-                    for event in neighbors[node_rid]:
-                        self.processor.handle(event=event, source=node_rid)
-                self.processor.flush_kobj_queue()
-                
+                self.poll()
                 elapsed = time.time() - start_time
                 sleep_time = self.config.koi_net.polling_interval - elapsed
                 if sleep_time > 0:

src/koi_net/server.py

@@ -109,20 +109,20 @@ async def poll_events(
         self, req: SignedEnvelope[PollEvents]
     ) -> SignedEnvelope[EventsPayload] | ErrorResponse:
         logger.info(f"Request to {POLL_EVENTS_PATH}")
-        events = self.event_queue.flush_poll_queue(req.payload.rid)
+        events = self.event_queue.flush_poll_queue(req.source_node)
         return EventsPayload(events=events)
 
     async def fetch_rids(
         self, req: SignedEnvelope[FetchRids]
     ) -> SignedEnvelope[RidsPayload] | ErrorResponse:
-        return self.response_handler.fetch_rids(req.payload)
+        return self.response_handler.fetch_rids(req.payload, req.source_node)
 
     async def fetch_manifests(
         self, req: SignedEnvelope[FetchManifests]
     ) -> SignedEnvelope[ManifestsPayload] | ErrorResponse:
-        return self.response_handler.fetch_manifests(req.payload)
+        return self.response_handler.fetch_manifests(req.payload, req.source_node)
 
     async def fetch_bundles(
         self, req: SignedEnvelope[FetchBundles]
     ) -> SignedEnvelope[BundlesPayload] | ErrorResponse:
-        return self.response_handler.fetch_bundles(req.payload)
+        return self.response_handler.fetch_bundles(req.payload, req.source_node)