<commit_before>import socket
import time
from binary_conversions import to_hex

UDP_BUFFER_SIZE = 65536
TCP_BUFFER_SIZE = 1000000

class _WithTimeouts(object):

    _default_timeout = 10

    def _get_timeout(self, timeout):
        if not timeout or str(timeout).lower() == 'none':
            return self._default_timeout
        elif str(timeout).lower() == 'blocking':
            return None
        return float(timeout)

    def _set_default_timeout(self, timeout):
        self._default_timeout = self._get_timeout(timeout)


class _Server(_WithTimeouts):

    def __init__(self, ip, port, timeout=None):
        self._ip = ip
        self._port = int(port)
        self._set_default_timeout(timeout)

    def _bind_socket(self):
        self._socket.bind((self._ip, self._port))
        self._is_connected = True

    def close(self):
        if self._is_connected:
            self._is_connected = False
            self._socket.close()
            self._message_stream = None

    def _get_message_stream(self, connection):
        if not self._protocol:
            return None
        return self._protocol.get_message_stream(BufferedStream(connection, self._default_timeout))


class UDPServer(_Server):

    def __init__(self, ip, port, timeout=None, protocol=None):
        _Server.__init__(self, ip, port, timeout)
        self._protocol = protocol
        self._last_client = None
        self._init_socket()

    def _init_socket(self):
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self._bind_socket()
        self._message_stream = self._get_message_stream(self)

    def receive_from(self, timeout=None, alias=None):
        self._check_no_alias(alias)
        timeout = self._get_timeout(timeout)
        self._socket.settimeout(timeout)
        msg, (ip, host) = self._socket.recvfrom(UDP_BUFFER_SIZE)
        print "*DEBUG* Read %s" % to_hex(msg)
        self._last_client = (ip, host)
        return msg, ip, host

    def _check_no_alias(self, alias):
        if alias:
            raise Exception('Connection aliases are not supported on UDP Servers')

    def receive(self, timeout=None, alias=None):
        return self.receive_from(timeout, alias)[0]

    def send_to(self, msg, ip, port):
        print "*DEBUG* Send %s" % to_hex(msg)
        self._socket.sendto(msg, (ip,int(port)))

    def send(self, msg, alias=None):
        if alias:
            raise Exception('UDP Server does not have connection aliases. Tried to use connection %s.' % alias)
        if not self._last_client:
            raise Exception('Server can not send to default client, because it has not received messages from clients.')
        self.send_to(msg, *self._last_client)

    def get_message(self, message_template, message_fields, timeout=None):
        msg = self._message_stream.get(message_template, timeout=timeout)
        errors = message_template.validate(msg, message_fields)
        if errors:
            print "Received %s" % repr(msg)
            print '\n'.join(errors)
            raise AssertionError(errors[0])
        print "*DEBUG* %s" % repr(msg)
        return msg


class TCPServer(_Server):

    def __init__(self, ip, port, timeout=None, protocol=None):
        _Server.__init__(self, ip, port, timeout)
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._bind_socket()
        self._socket.listen(1)
        self._connections = _NamedCache('connection')
        self._protocol = protocol
        self._message_stream = self._get_message_stream(self)

    def receive(self, timeout=None, alias=None):
        return self.receive_from(timeout, alias)[0]

    def receive_from(self, timeout=None, alias=None):
        connection, (ip, host) = self._connections.get(alias)
        timeout = self._get_timeout(timeout)
        connection.settimeout(timeout)
        msg = connection.recv(TCP_BUFFER_SIZE)
        print "*DEBUG* Read %s" % to_hex(msg)
        return msg, ip, host

    def accept_connection(self, alias=None):
        connection, client_address = self._socket.accept()
        self._connections.add((connection, client_address), alias)
        return client_address

    def send(self, msg, alias=None):
        print "*DEBUG* Send %s" % to_hex(msg)
        self._connections.get(alias)[0].send(msg)

    def send_to(self, *args):
        raise Exception("TCP server cannot send to a specific address.")

    def close(self):
        if self._is_connected:
            self._is_connected = False
            for connection, _ in self._connections:
                connection.close()
            self._socket.close()

    def get_message(self, message_template, message_fields, timeout=None, alias=None):
        msg = self._message_stream.get(message_template, timeout=timeout)
        # TODO: duplication with UDPServer and _Client
        # TODO: Wrap connection to server like wrapper with message logging
        errors = message_template.validate(msg, message_fields)
        if errors:
            print "Received %s" % repr(msg)
            print '\n'.join(errors)
            raise AssertionError(errors[0])
        print "*DEBUG* %s" % repr(msg)
        return msg
<commit_msg>move get_message in tcp and udp servers to superclass<commit_after>import socket
import time
from binary_conversions import to_hex

UDP_BUFFER_SIZE = 65536
TCP_BUFFER_SIZE = 1000000

class _WithTimeouts(object):

    _default_timeout = 10

    def _get_timeout(self, timeout):
        if not timeout or str(timeout).lower() == 'none':
            return self._default_timeout
        elif str(timeout).lower() == 'blocking':
            return None
        return float(timeout)

    def _set_default_timeout(self, timeout):
        self._default_timeout = self._get_timeout(timeout)


class _Server(_WithTimeouts):

    def __init__(self, ip, port, timeout=None):
        self._ip = ip
        self._port = int(port)
        self._set_default_timeout(timeout)

    def _bind_socket(self):
        self._socket.bind((self._ip, self._port))
        self._is_connected = True

    def close(self):
        if self._is_connected:
            self._is_connected = False
            self._socket.close()
            self._message_stream = None

    def _get_message_stream(self, connection):
        if not self._protocol:
            return None
        return self._protocol.get_message_stream(BufferedStream(connection, self._default_timeout))

    def get_message(self, message_template, message_fields, timeout=None, alias=None):
        # TODO: duplication with UDPServer and _Client
        # TODO: Wrap connection to server like wrapper with message logging
        msg = self._message_stream.get(message_template, timeout=timeout)
        errors = message_template.validate(msg, message_fields)
        if errors:
            print "Received %s" % repr(msg)
            print '\n'.join(errors)
            raise AssertionError(errors[0])
            print "*DEBUG* %s" % repr(msg)
        return msg


class UDPServer(_Server):

    def __init__(self, ip, port, timeout=None, protocol=None):
        _Server.__init__(self, ip, port, timeout)
        self._protocol = protocol
        self._last_client = None
        self._init_socket()

    def _init_socket(self):
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self._bind_socket()
        self._message_stream = self._get_message_stream(self)

    def receive_from(self, timeout=None, alias=None):
        self._check_no_alias(alias)
        timeout = self._get_timeout(timeout)
        self._socket.settimeout(timeout)
        msg, (ip, host) = self._socket.recvfrom(UDP_BUFFER_SIZE)
        print "*DEBUG* Read %s" % to_hex(msg)
        self._last_client = (ip, host)
        return msg, ip, host

    def _check_no_alias(self, alias):
        if alias:
            raise Exception('Connection aliases are not supported on UDP Servers')

    def receive(self, timeout=None, alias=None):
        return self.receive_from(timeout, alias)[0]

    def send_to(self, msg, ip, port):
        print "*DEBUG* Send %s" % to_hex(msg)
        self._socket.sendto(msg, (ip,int(port)))

    def send(self, msg, alias=None):
        if alias:
            raise Exception('UDP Server does not have connection aliases. Tried to use connection %s.' % alias)
        if not self._last_client:
            raise Exception('Server can not send to default client, because it has not received messages from clients.')
        self.send_to(msg, *self._last_client)


class TCPServer(_Server):

    def __init__(self, ip, port, timeout=None, protocol=None):
        _Server.__init__(self, ip, port, timeout)
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._bind_socket()
        self._socket.listen(1)
        self._connections = _NamedCache('connection')
        self._protocol = protocol
        self._message_stream = self._get_message_stream(self)

    def receive(self, timeout=None, alias=None):
        return self.receive_from(timeout, alias)[0]

    def receive_from(self, timeout=None, alias=None):
        connection, (ip, host) = self._connections.get(alias)
        timeout = self._get_timeout(timeout)
        connection.settimeout(timeout)
        msg = connection.recv(TCP_BUFFER_SIZE)
        print "*DEBUG* Read %s" % to_hex(msg)
        return msg, ip, host

    def accept_connection(self, alias=None):
        connection, client_address = self._socket.accept()
        self._connections.add((connection, client_address), alias)
        return client_address

    def send(self, msg, alias=None):
        print "*DEBUG* Send %s" % to_hex(msg)
        self._connections.get(alias)[0].send(msg)

    def send_to(self, *args):
        raise Exception("TCP server cannot send to a specific address.")

    def close(self):
        if self._is_connected:
            self._is_connected = False
            for connection, _ in self._connections:
                connection.close()
            self._socket.close()

{'fraction_non_alphanumeric': 0.10030424968102856, 'fraction_numerical': 0.0040239473942487, 'mean_word_length': 2.779673590504451, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4366851', 'n_tokens_mistral': 2992, 'n_tokens_neox': 2882, 'n_words': 725}