controller-wisetopdataserver/python/biopro/util/socket.py

import abc
import pickle as _pickle
import struct as _struct
from functools import wraps
from socket import socket as Socket, timeout as _timeout, AF_UNIX, SOCK_STREAM
from typing import Any, Tuple, Optional, Type, Dict

from biopro.text import *

SocketTimeoutError = _timeout


def socket_send(cnt: Socket, data: Any, addr: Optional[Tuple[str, int]] = None) -> None:
    """send data through the socket.

    :param cnt: socket connection
    :param data:
    :param addr:
    """
        
    buff = _pickle.dumps(data)
    send = _struct.pack('>I', len(buff)) + buff
    if addr is not None:
        cnt.sendto(send, addr)
    else:
        try:
            cnt.sendall(send)
        except:
            raise RuntimeError('socket send fault')


def socket_recv(cnt: Socket, timeout: Optional[float] = 0.0) -> Any:
    """receive data through the socket

    :param cnt: socket connection
    :param timeout:
    :return: data
    """
    cnt.settimeout(timeout)
    try:
        ret = cnt.recv(4)
    except ConnectionResetError as e:
        return None

    if ret is None or len(ret) == 0:
        return None

    data_length = _struct.unpack('>I', ret)[0]

    # data_length overflow
    if int(data_length) >= 0x80000000:
        return None

    data = b''
    while len(data) < data_length:
        ret = cnt.recv(data_length - len(data))

        if ret is None:
            return None

        data += ret

    ret = _pickle.loads(data)

    if isinstance(ret, RuntimeError):
        print('RuntimeError', ret)
        return None
        # raise ret

    return ret


class SocketClient(metaclass=abc.ABCMeta):
    __slots__ = ('_socket_file', '_socket')

    def __init__(self, socket_file: str):
        self._socket_file = socket_file
        self._socket: Optional[Socket] = None

    def __enter__(self):
        self.open_socket()

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close_socket()

    @property
    def socket_opened(self) -> bool:
        return self._socket is not None

    def open_socket(self):
        if self._socket is None:
            try:
                self._socket = Socket(AF_UNIX, SOCK_STREAM)
                self._socket.settimeout(5)
                self._socket.connect(self._socket_file)
            except Exception as e:
                print('open_socket error:', e)
            # except FileNotFoundError as e:
            #     raise RuntimeError('server not available') from e

            # except ConnectionRefusedError as e:
            #     raise RuntimeError('server crash') from e

    def close_socket(self):
        if self._socket is not None:
            self._socket.close()

        self._socket = None

    def send_command(self, cmd: str, *args: Any, **option) -> Any:
        if self._socket is None:
            raise RuntimeError(SOCKET_NOT_OPEN)

        socket_send(self._socket, (cmd, args, option))

        return socket_recv(self._socket, timeout=None)


# noinspection PyPep8Naming
def SocketClientMacro(api_class: Type):
    """**meta-class** generating the lost implemented method defined in *api_class* for SocketClient subclass.

    code example ::

        # noinspection PyAbstractClass
        class ApiClient(SocketClient, ApiClass, metaclass=SocketClientMacro(ApiClass)):
            pass

    default implement method look-like ::

        def method_not_implement(self, *argv, **options):
            return self._send_command('method_not_implement', *argv, **options)

        method_not_implement.__ref_func = api_class.method_not_implement
        method_not_implement.__doc__ = "(auto-generated)"

    The IDE would warning that the SocketClient subclass do not implement all abstract method in *api_class*.
    To suppressed that, add ``noinspection PyAbstractClass`` comment above class definition for IDEA.

    :param api_class: api class reference
    :return: meta function
    """

    def _meta(define_class_name: str, super_class: Tuple[Type, ...], class_attr: Dict[str, Any]):
        if len(super_class) == 0 or super_class[0] != SocketClient:
            raise RuntimeError(define_class_name + ' not a subclass of SocketClient')

        elif api_class not in super_class:
            raise RuntimeError(define_class_name + ' not a subclass of ' + api_class.__name__)

        for api_attr_name in dir(api_class):

            # should be public
            if not api_attr_name.startswith('_'):
                api_func = getattr(api_class, api_attr_name)

                # should be a method and not overwrote
                if callable(api_func) and api_attr_name not in class_attr:
                    # get callable by calling function to create local variable binding.
                    class_attr[api_attr_name] = _gen_impl_func(api_func)

        return type(define_class_name, super_class, class_attr)

    return _meta


def _gen_impl_func(ref_func):
    @wraps(ref_func)
    def _impl_func(self: SocketClient, *argv, **options):
        return self.send_command(ref_func.__name__, *argv, **options)

    # additional attribute, make django can get the actual function information
    # instead the information of function _impl_func

    _impl_func.__ref_func = ref_func
    _impl_func.__isabstractmethod__ = False
    _impl_func.__doc__ = "(auto-generated)"

    return _impl_func