Add types

re6st/registry.py

@@ -201,10 +201,10 @@ class RegistryServer:
         x = utils.packInteger(1 + utils.unpackInteger(self.version)[0])
         self.version = x + self.cert.sign(x)
 
-    def sendto(self, prefix, code):
+    def sendto(self, prefix: str, code: int):
         self.sock.sendto(("%s\0%c" % (prefix, code)).encode("utf-8"), ('::1', tunnel.PORT))
 
-    def recv(self, code):
+    def recv(self, code: int) -> (str, str):
         try:
             data = self.sock.recv(1<<16)
             logging.info("recv raw: %r", data)
@@ -215,7 +215,7 @@ class RegistryServer:
         else:
             if msg:
                 if msg[0:1] == bytes([code]):
-                    return prefix.decode(), msg[1:]
+                    return prefix.decode(), msg[1:].decode()
                 else:
                     logging.error("Unexpected code: %r", msg)
             else:
@@ -615,7 +615,7 @@ class RegistryServer:
                     v and base64.b64encode(x509.encrypt(cert, v)).decode("ascii")
         return zlib.compress(json.dumps(config).encode("utf-8"))
 
-    def _queryAddress(self, peer):
+    def _queryAddress(self, peer) -> str:
         logging.info("Querying address for %s/%s %r", int(peer, 2), len(peer), peer)
         self.sendto(peer, 1)
         s = self.sock,
@@ -626,13 +626,13 @@ class RegistryServer:
             prefix, msg = self.recv(1)
             logging.info("* received: %r - %r", prefix, msg)
             if prefix == peer:
-                return msg.decode()
+                return msg
             timeout = max(0, end - time.time())
         logging.info("Timeout while querying address for %s/%s",
                      int(peer, 2), len(peer))
 
     @rpc
-    def getCountry(self, cn, address):
+    def getCountry(self, cn, address) -> bytes:
         country = self._geoiplookup(address)[0]
         return None if country == '*' else country.encode()