add new abc for sockets

[erebus.git] / modlib.py

diff --git a/modlib.py b/modlib.py
index b36c153287f4979c107a8768c60b91536949ec15..11ca1fb7a4acc454cab79b2254583fe8d1d885d8 100644 (file)
--- a/modlib.py
+++ b/modlib.py
@@ -3,6 +3,7 @@
 # module helper functions, see modules/modtest.py for usage
 # This file is released into the public domain; see http://unlicense.org/
 
+import abc
 import sys
 import socket
 from functools import wraps
@@ -60,6 +61,7 @@ class modlib(object):
        WRONGARGS = "Wrong number of arguments."
 
        def __init__(self, name):
+               self.Socketlike = Socketlike
                self.hooks = {} # {command:handler}
                self.chanhooks = {} # {channel:handler}
                self.exceptionhooks = [] # [(exception,handler)]
@@ -206,9 +208,7 @@ class modlib(object):
                return self._hooksocket(socket.AF_UNIX, socket.SOCK_STREAM, path, data)
        def _hooksocket(self, af, ty, address, data):
                def realhook(cls):
-                       if not (hasattr(cls, 'getdata') and callable(cls.getdata)):
-                               # Check early that the object implements getdata.
-                               # If getdata ever returns a non-empty list, then a parse method must also exist, but we don't check that.
+                       if not issubclass(cls, Socketlike):
                                raise Exception('Attempted to hook a socket without a class to process data')
                        self.sockhooks.append((af, ty, address, cls, data))
                        if self.parent is not None:
@@ -322,7 +322,52 @@ class modlib(object):
                        return func
                return realhook
 
-class _ListenSocket(object):
+class Socketlike(abc.ABC):
+       def __init__(self, sock, data):
+               """This default method saves the socket in self.sock and creates self.buffer for getdata(). The data is discarded."""
+               self.sock = sock
+               self.buffer = b''
+
+       def getdata(self):
+               """This default method gets LF or CRLF separated lines from the socket and returns an array of completely-seen lines to the core.
+               This should work well for most line-based protocols (like IRC)."""
+               recvd = self.sock.recv(8192)
+               if recvd == b"": # EOF
+                       if len(self.buffer) != 0:
+                               # Process what's left in the buffer. We'll get called again after.
+                               remaining_buf = self.buffer.decode('utf-8', 'backslashreplace')
+                               self.buffer = b""
+                               return [remaining_buf]
+                       else:
+                               # Nothing left in the buffer. Return None to signal the core to close this socket.
+                               return None
+               self.buffer += recvd
+               lines = []
+
+               while b"\n" in self.buffer:
+                       pieces = self.buffer.split(b"\n", 1)
+                       s = pieces[0].decode('utf-8', 'backslashreplace').rstrip("\r")
+                       lines.append(pieces[0].decode('utf-8', 'backslashreplace'))
+                       self.buffer = pieces[1]
+
+               return lines
+
+       def __str__(self):
+               return '%s#%d' % (self.__class__.__name__, self.sock.fileno())
+       def __repr__(self):
+               return '<%s.%s #%d %s:%d>' % ((self.__class__.__module__, self.__class__.__name__, self.sock.fileno())+self.sock.getpeername())
+
+       @abc.abstractmethod
+       def parse(self, chunk): pass
+
+       @classmethod
+       def __subclasshook__(cls, C):
+               if cls is Socketlike:
+                       if any('parse' in B.__dict__ and 'getdata' in B.__dict__ for B in C.__mro__):
+                               return True
+               return NotImplemented
+
+class _ListenSocket(Socketlike):
        def __init__(self, lib, sock, cls, data):
                self.clients = []
                self.lib = lib
@@ -343,6 +388,10 @@ class _ListenSocket(object):
                        self.clients.remove((client,obj))
                return close
 
+       def parse(self):
+               # getdata will never return a non-empty array, so parse will never be called; but Socketlike requires this method
+               pass
+
        def getdata(self):
                client, addr = self.sock.accept()
                obj = self.cls(client, self.data)