supwisdom/protocol/nettrans.py

# -*- coding: utf-8 -*-
""" 一卡通后台通讯协议包封装
所有数据都是 unicode  utf-8 编码
"""
import traceback
import simplejson as json


def _unicode_str(value):
    if isinstance(value, str):
        return value.decode('utf-8')
    elif isinstance(value, unicode):
        return value
    else:
        return unicode(value)


class message_writer:
    """ 后台消息写接口
    """
    def __init__(self):
        self.clear()

    def clear(self):
        self._attr = {}
        self._row_data = []
        self._col_names = []
        self._row = {}
        self._frozen = False

    def attr(self, name, value):
        if name in self._attr:
            raise ValueError('Attribute [%s] already exists' % name)
        self._attr[name] = _unicode_str(value)
        return self

    def row(self, name, value):
        if not (name in self._col_names):
            self._col_names.append(name)
        self._row[name] = _unicode_str(value)

    def add_row(self):
        if not self._row:
            raise ValueError('Row has no values')
        self._row_data.append(self._row)
        self._row = {}
        return len(self._row_data)

    def append_row(self, row):
        if not self._col_names:
            for k, v in row.iteritems():
                self._col_names.append(k)
            self._row_data.append(row)
        else:
            new_row = dict([(v, row.get(v, None)) for v in self._col_names])
            self._row_data.append(new_row)

    def _normalize(self):
        colcnt = len(self._col_names)
        rowcnt = len(self._row_data)
        self._attr["colcnt"] = colcnt
        self._attr["rowcnt"] = rowcnt
        if colcnt > 0:
            self._attr["colname"] = self._col_names
            self._attr["coldesc"] = self._col_names

        if rowcnt > 0 and colcnt <= 0:
            raise ValueError(u'数据没有列值')

        if rowcnt > 0:
            self._attr["rowdata"] = [[row.get(colname, u"") for colname in self._col_names]
                for row in self._row_data]
        self._frozen = True

    def serialize(self):
        """ 返回 unicode utf-8编码
        """
        if not self._frozen:
            self._normalize()
        seri = json.dumps(self._attr, ensure_ascii=False, encoding='utf-8')
        return seri

    def root(self):
        if not self._frozen:
            self._normalize()
        return self._attr


class message_reader:
    """ 后台消息读接口
    """
    def __init__(self):
        self.clear()

    def clear(self):
        """ 清空数据
        """
        self._attr = {}
        self._row_data = None
        self._col_names = []
        self._col_desc = []
        self._row_no = -1
        self._frozen = False

    def unserialize(self, message):
        """ 解析报文
        """
        self.clear()
        try:
            message = _unicode_str(message)
            obj = json.loads(message)
            #obj = json.loads(msg.replace('\\',"\\\\"))
        except Exception, ex:
            traceback.print_exc()
            raise ValueError(u'解析报文错误,' + ex.message)

        if u"colname" in obj:
            colname = obj[u"colname"]
            colcnt = obj[u"colcnt"]
            rowcnt = obj[u"rowcnt"]

            if colcnt != len(colname):
                raise ValueError(u'记录列数定义不符')

            rowdata = obj[u"rowdata"]
            if rowcnt != len(rowdata):
                raise ValueError(u'记录行数定义不符')

            self._row_data = [dict(zip(colname, row)) for row in rowdata]
            self._col_names = colname
            self._col_desc = obj["coldesc"]
        else:
            pass

        for n, v in obj.iteritems():
            if n in (u"colname", u"colcnt", u"rowcnt", u"coldesc", u"rowdata"):
                continue
            self._attr[n] = v
        self._frozen = True
        #print self._attr

    def attr(self, name):
        """ 获取属性值，如果没有改属性返回空字符串
        """
        assert self._frozen == True
        if name not in self._attr:
            return ""
        return self._attr[name]

    def has_more_row(self):
        """ 判断是否有后续行
        """
        assert self._frozen == True
        if not self._row_data:
            return False
        if (self._row_no + 1) < len(self._row_data):
            return True
        return False

    def next_row(self):
        """ 将读取指针移到下一行
        """
        assert self._frozen == True
        assert self.has_more_row()
        self._row_no += 1

    _ARG_DEFAULT = []

    def col(self, name, default=[], strip=True):
        """ 获取当前行的字段值
        """
        assert self._row_no >= 0
        arg = self._row_data[self._row_no].get(name, None)
        if arg is None:
            if strip and default == message_reader._ARG_DEFAULT:
                raise IndexError(u'no such column %s' % name)
            return default
        return arg

    def is_eof(self):
        """判断是否到记录尾
        """
        if self._row_no >= len(self._row_data):
            return True
        return False

    def fetch_rows(self):
        """ 读取所有行记录
        """
        if not self.has_more_row():
            return
        self._row_no = 0
        for row in self._row_data:
            self._row_no += 1
            yield row