Raft_Shellmiao/network_api.py

import json
import socket


class NetworkAPI(object):
    def __init__(self, addr=None, timeout=None):
        self.addr = tuple(addr)
        self.ss = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        if addr:
            self.bind(tuple(addr))
        if timeout:
            self.ss.settimeout(timeout)

    def bind(self, addr):
        self.addr = tuple(addr)
        self.ss.bind(addr)

    def set_timeout(self, timeout):
        self.ss.settimeout(timeout)

    def send(self, data, addr):
        data = json.dumps(data).encode("utf-8")
        self.ss.sendto(data, tuple(addr))

    def receive(self, addr=None, timeout=None):
        if addr:
            self.bind(addr)
        if not self.addr:
            raise "please bind to an address"

        if timeout:
            self.set_timeout(timeout)

        data, addr = self.ss.recvfrom(65535)
        return json.loads(data), addr

    def close(self):
        self.ss.close()