vstf/vstf/common/message.py

   1 import json
   2 import uuid
   3 import logging
   4 import traceback
   5 from vstf.common import constants
   6
   7 LOG = logging.getLogger(__name__)
   8
   9
  10 def json_defaults(obj):
  11     if isinstance(obj, set):
  12         return list(obj)
  13     return "unknow obj"
  14
  15
  16 def encode(msg):
  17     """obj to string"""
  18     if isinstance(msg, str):
  19         return msg
  20     else:
  21         return json.dumps(msg, default=json_defaults)
  22
  23
  24 def decode(msg):
  25     """string to obj"""
  26     if isinstance(msg, str):
  27         return json.loads(msg)
  28     else:
  29         return msg
  30
  31
  32 def gen_corrid():
  33     return str(uuid.uuid4())
  34
  35
  36 def add_context(msg, **kwargs):
  37     return {'head': kwargs, 'body': msg}
  38
  39
  40 def get_context(msg):
  41     if "head" in msg.iterkeys():
  42         return msg['head']
  43     else:
  44         return ""
  45
  46
  47 def get_body(msg):
  48     if "body" in msg.iterkeys():
  49         return msg['body']
  50     else:
  51         return None
  52
  53
  54 def get_corrid(context):
  55     """
  56     :param return: string of corrid or empty
  57     """
  58     if "corrid" in context.iterkeys():
  59         return context['corrid']
  60     else:
  61         return ""
  62
  63
  64 def send(func, data):
  65     # the message must be a string
  66     if not isinstance(data, str):
  67         raise ValueError("the data must be a string")
  68
  69     # the message's len must > 0
  70     msg_len = len(data)
  71     if msg_len <= 0:
  72         return True
  73
  74     # the message's len must be less 999999999
  75     if len(str(msg_len)) > constants.MSG_FLAG_LEN:
  76         raise ValueError("the data's len too long")
  77
  78     data = (constants.MSG_FLAG % (msg_len)) + data
  79     total_send = msg_len + constants.MSG_FLAG_LEN
  80
  81     count = 0
  82     while count < total_send:
  83         sent = func(data[count:])
  84         if 0 == sent:
  85             raise RuntimeError("socket connection broken")
  86         count += sent
  87
  88     return msg_len
  89
  90
  91 def sendto(func, data, addr):
  92     # the message must be a string
  93     if not isinstance(data, str):
  94         raise ValueError("the data must be a string")
  95
  96     # the message's len must > 0
  97     msg_len = len(data)
  98     if msg_len <= 0:
  99         return True
 100
 101     # the message's len must be less 999999999
 102     if len(str(msg_len)) > constants.MSG_FLAG_LEN:
 103         raise ValueError("the data's len too long")
 104
 105     data = (constants.MSG_FLAG % (msg_len)) + data
 106     total_send = msg_len + constants.MSG_FLAG_LEN
 107
 108     count = 0
 109     while count < total_send:
 110         sent = func(data[count:], addr)
 111         if 0 == sent:
 112             raise RuntimeError("socket connection broken")
 113         count += sent
 114
 115     return msg_len
 116
 117
 118 def recv(func):
 119     head = func(constants.MSG_FLAG_LEN)
 120     # the FIN change to '' in python
 121     if head == '':
 122         raise RuntimeError("socket connection broken")
 123
 124     if not head.isdigit():
 125         raise ValueError("the msg head is not a num.")
 126
 127     msg_len = int(head)
 128     chunks = []
 129     count = 0
 130     while count < msg_len:
 131         chunk = func(min(msg_len - count, constants.buff_size))
 132         if chunk == '':
 133             raise RuntimeError("socket connection broken")
 134         chunks.append(chunk)
 135         count += len(chunk)
 136
 137     return ''.join(chunks)
 138
 139
 140 def dumpstrace():
 141     return traceback.format_exc()