From b2cfd2f036414650a9ca841b03220b23c3126e6a Mon Sep 17 00:00:00 2001
From: thetestgame <nxtlabs13@gmail.com>
Date: Mon, 11 Nov 2024 00:16:46 -0600
Subject: [PATCH] Cleaned up control messages

---
 panda3d_astron/interfaces/state.py |  1 -
 panda3d_astron/internal.py         | 24 ++++++++++++++++++------
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/panda3d_astron/interfaces/state.py b/panda3d_astron/interfaces/state.py
index f5a025a..4cef942 100644
--- a/panda3d_astron/interfaces/state.py
+++ b/panda3d_astron/interfaces/state.py
@@ -545,7 +545,6 @@ def request_delete(self, do: object) -> None:
     requestDelete = request_delete
 
 
-
     def set_owner(self, doId: int, newOwner: int) -> None:
         """
         Sets the owner of a DistributedObject. This will enable the new owner to send "ownsend" fields,
diff --git a/panda3d_astron/internal.py b/panda3d_astron/internal.py
index 9554823..f0aaca0 100644
--- a/panda3d_astron/internal.py
+++ b/panda3d_astron/internal.py
@@ -115,7 +115,7 @@ def get_context(self) -> int:
         self.__contextCounter = (self.__contextCounter + 1) & 0xFFFFFFFF
         return self.__contextCounter
     
-    def allocate_channel(self):
+    def allocate_channel(self) -> None:
         """
         Allocate an unused channel out of this AIR's configured channel space.
         This is also used to allocate IDs for DistributedObjects, since those
@@ -126,7 +126,7 @@ def allocate_channel(self):
     
     allocateChannel = allocate_channel
 
-    def deallocate_channel(self, channel):
+    def deallocate_channel(self, channel: int) -> None:
         """
         Return the previously-allocated channel back to the allocation pool.
         """
@@ -135,7 +135,7 @@ def deallocate_channel(self, channel):
 
     deallocateChannel = deallocate_channel
 
-    def register_for_channel(self, channel):
+    def register_for_channel(self, channel: int) -> None:
         """
         Register for messages on a specific Message Director channel.
         If the channel is already open by this AIR, nothing will happen.
@@ -152,7 +152,18 @@ def register_for_channel(self, channel):
 
     registerForChannel = register_for_channel
 
-    def unregister_for_channel(self, channel):
+    def register_for_location_channel(self, parentId: int, zoneId: int) -> None:
+        """
+        Register for messages on a specific state server zone channel.
+        If the channel is already open by this AIR, nothing will happen.
+        """
+
+        channel = (parentId << 32) | zoneId
+        self.registerForChannel(channel)
+
+    registerForLocationChannel = register_for_location_channel
+
+    def unregister_for_channel(self, channel: int) -> None:
         """
         Unregister a channel subscription on the Message Director. The Message
         Director will cease to relay messages to this AIR sent on the channel.
@@ -169,7 +180,7 @@ def unregister_for_channel(self, channel):
 
     unregisterForChannel = unregister_for_channel
 
-    def add_post_remove(self, dg):
+    def add_post_remove(self, dg: object) -> None:
         """
         Register a datagram with the Message Director that gets sent out if the
         connection is ever lost.
@@ -182,11 +193,12 @@ def add_post_remove(self, dg):
         dg2.addServerControlHeader(msgtypes.CONTROL_ADD_POST_REMOVE)
         dg2.add_uint64(self.ourChannel)
         dg2.add_blob(dg.getMessage())
+
         self.send(dg2)
 
     addPostRemove = add_post_remove
 
-    def clear_post_remove(self):
+    def clear_post_remove(self) -> None:
         """
         Clear all datagrams registered with addPostRemove.
         This is useful if the Panda3D process is performing a clean exit. It may