diff --git a/.gitignore b/.gitignore index 1f0a52be..969df946 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,6 @@ build/* dist/* -test/* +tests/* itchat.egg-info/* *.pyc *.swp diff --git a/itchat/components/contact.py b/itchat/components/contact.py index 35793d64..ea7ff00b 100644 --- a/itchat/components/contact.py +++ b/itchat/components/contact.py @@ -6,7 +6,7 @@ from .. import config, utils from ..returnvalues import ReturnValue -from ..storage import contact_change +from ..storage import contact_change, templates logger = logging.getLogger('itchat') @@ -99,8 +99,8 @@ def update_friend(self, userName): return r if len(r) != 1 else r[0] def update_info_dict(oldInfoDict, newInfoDict): - ''' - only normal values will be updated here + ''' only normal values will be updated here + because newInfoDict is normal dict, so it's not necessary to consider templates ''' for k, v in newInfoDict.items(): if any((isinstance(v, t) for t in (tuple, list, dict))): @@ -126,8 +126,8 @@ def update_local_chatrooms(core, l): if oldChatroom: update_info_dict(oldChatroom, chatroom) # - update other values - memberList, oldMemberList = (c.get('MemberList', []) - for c in (chatroom, oldChatroom)) + memberList = chatroom.get('MemberList', []) + oldMemberList = oldChatroom.memberList if memberList: for member in memberList: oldMember = utils.search_dict_list( @@ -137,17 +137,19 @@ def update_local_chatrooms(core, l): else: oldMemberList.append(member) else: - oldChatroom = chatroom - core.chatroomList.append(chatroom) + oldChatroom = templates.wrap_user_dict(chatroom) + core.chatroomList.append(oldChatroom) # delete useless members if len(chatroom['MemberList']) != len(oldChatroom['MemberList']) and \ chatroom['MemberList']: existsUserNames = [member['UserName'] for member in chatroom['MemberList']] delList = [] for i, member in enumerate(oldChatroom['MemberList']): - if member['UserName'] not in existsUserNames: delList.append(i) + if member['UserName'] not in existsUserNames: + delList.append(i) delList.sort(reverse=True) - for i in delList: del oldChatroom['MemberList'][i] + for i in delList: + del oldChatroom['MemberList'][i] # - update OwnerUin if oldChatroom.get('ChatRoomOwner') and oldChatroom.get('MemberList'): oldChatroom['OwnerUin'] = utils.search_dict_list(oldChatroom['MemberList'], diff --git a/itchat/components/login.py b/itchat/components/login.py index c35d1ecd..6b5fb043 100644 --- a/itchat/components/login.py +++ b/itchat/components/login.py @@ -9,6 +9,7 @@ from .. import config, utils from ..returnvalues import ReturnValue +from ..storage.templates import wrap_user_dict from .contact import update_local_chatrooms, update_local_friends from .messages import produce_msg @@ -182,7 +183,7 @@ def web_init(self): # deal with login info utils.emoji_formatter(dic['User'], 'NickName') self.loginInfo['InviteStartCount'] = int(dic['InviteStartCount']) - self.loginInfo['User'] = utils.struct_friend_info(dic['User']) + self.loginInfo['User'] = wrap_user_dict(utils.struct_friend_info(dic['User'])) self.memberList.append(self.loginInfo['User']) self.loginInfo['SyncKey'] = dic['SyncKey'] self.loginInfo['synckey'] = '|'.join(['%s_%s' % (item['Key'], item['Val']) diff --git a/itchat/core.py b/itchat/core.py index dea80903..95ba1d52 100644 --- a/itchat/core.py +++ b/itchat/core.py @@ -18,7 +18,7 @@ def __init__(self): - failing is failing ''' self.alive, self.isLogging = False, False - self.storageClass = storage.Storage() + self.storageClass = storage.Storage(self) self.memberList = self.storageClass.memberList self.mpList = self.storageClass.mpList self.chatroomList = self.storageClass.chatroomList diff --git a/itchat/returnvalues.py b/itchat/returnvalues.py index e9643d08..b9600bd6 100644 --- a/itchat/returnvalues.py +++ b/itchat/returnvalues.py @@ -4,6 +4,23 @@ TRANSLATE = 'Chinese' class ReturnValue(dict): + ''' turn return value of itchat into a boolean value + for requests: + ..code::python + + import requests + r = requests.get('http://httpbin.org/get') + print(ReturnValue(rawResponse=r) + + for normal dict: + ..code::python + + returnDict = { + 'BaseResponse': { + 'Ret': 0, + 'ErrMsg': 'My error msg', }, } + print(ReturnValue(returnDict)) + ''' def __init__(self, returnValueDict={}, rawResponse=None): if rawResponse: try: @@ -14,7 +31,8 @@ def __init__(self, returnValueDict={}, rawResponse=None): 'Ret': -1004, 'ErrMsg': 'Unexpected return value', }, 'Data': rawResponse.content, } - for k, v in returnValueDict.items(): self[k] = v + for k, v in returnValueDict.items(): + self[k] = v if not 'BaseResponse' in self: self['BaseResponse'] = { 'ErrMsg': 'no BaseResponse in raw response', @@ -45,6 +63,7 @@ def __repr__(self): -1003: u'服务器拒绝连接', -1004: u'服务器返回异常值', -1005: u'参数错误', + -1006: u'无效操作', 0: u'请求成功', }, } diff --git a/itchat/storage.py b/itchat/storage/__init__.py similarity index 79% rename from itchat/storage.py rename to itchat/storage/__init__.py index 5bfff6f0..22419d07 100644 --- a/itchat/storage.py +++ b/itchat/storage/__init__.py @@ -5,6 +5,10 @@ import queue as Queue from threading import Lock +from .templates import ( + ContactList, AbstractUserDict, User, + MassivePlatform, Chatroom, ChatroomMember) + def contact_change(fn): def _contact_change(core, *args, **kwargs): with core.storageClass.updateLock: @@ -12,15 +16,21 @@ def _contact_change(core, *args, **kwargs): return _contact_change class Storage(object): - def __init__(self): + def __init__(self, core): self.userName = None self.nickName = None self.updateLock = Lock() - self.memberList = [] - self.mpList = [] - self.chatroomList = [] + self.memberList = ContactList() + self.mpList = ContactList() + self.chatroomList = ContactList() self.msgList = Queue.Queue(-1) self.lastInputUserName = None + self.memberList.set_default_value(contactClass=User) + self.memberList.core = core + self.mpList.set_default_value(contactClass=MassivePlatform) + self.mpList.core = core + self.chatroomList.set_default_value(contactClass=Chatroom) + self.chatroomList.core = core def dumps(self): return { 'userName' : self.userName, @@ -33,11 +43,14 @@ def loads(self, j): self.userName = j.get('userName', None) self.nickName = j.get('nickName', None) del self.memberList[:] - for i in j.get('memberList', []): self.memberList.append(i) + for i in j.get('memberList', []): + self.memberList.append(i) del self.mpList[:] - for i in j.get('mpList', []): self.mpList.append(i) + for i in j.get('mpList', []): + self.mpList.append(i) del self.chatroomList[:] - for i in j.get('chatroomList', []): self.chatroomList.append(i) + for i in j.get('chatroomList', []): + self.chatroomList.append(i) self.lastInputUserName = j.get('lastInputUserName', None) def search_friends(self, name=None, userName=None, remarkName=None, nickName=None, wechatAccount=None): diff --git a/itchat/storage/templates.py b/itchat/storage/templates.py new file mode 100644 index 00000000..e1430e74 --- /dev/null +++ b/itchat/storage/templates.py @@ -0,0 +1,244 @@ +import logging, copy, pickle + +from ..returnvalues import ReturnValue + +logger = logging.getLogger('itchat') + +class UnInitializedItchat(object): + def _raise_error(self, *args, **kwargs): + logger.warning('An itchat instance is called before initialized') + def __getattr__(self, value): + return self._raise_error + +fakeItchat = UnInitializedItchat() + +class ContactList(list): + ''' when a dict is append, init function will be called to format that dict + ''' + def __init__(self, *args, **kwargs): + super(ContactList, self).__init__(*args, **kwargs) + self.contactInitFn = None + self.contactClass = User + self.core = fakeItchat + def set_default_value(self, initFunction=None, contactClass=None): + if hasattr(initFunction, '__call__'): + self.contactInitFn = initFunction + if hasattr(contactClass, '__call__'): + self.contactClass = contactClass + def append(self, value): + contact = self.contactClass(value) + contact.core = self.core + if self.contactInitFn is not None: + contact = self.contactInitFn(contact) + super(ContactList, self).append(contact) + def __deepcopy__(self, memo): + return self.__class__([copy.deepcopy(v) for v in self]) + def __getstate__(self): + return [pickle.dumps(v) for v in self] + def __setstate__(self, state): + for v in state: + super(ContactList, self).append(pickle.loads(v)) + +fakeContactList = ContactList + +class AbstractUserDict(dict): + def __init__(self, *args, **kwargs): + super(AbstractUserDict, self).__init__(*args, **kwargs) + self.core = fakeItchat + def update(self): + return ReturnValue({'BaseResponse': { + 'Ret': -1006, + 'ErrMsg': '%s can not be updated' % \ + self.__class__.__name__, }, }) + def set_alias(self, alias): + return ReturnValue({'BaseResponse': { + 'Ret': -1006, + 'ErrMsg': '%s can not set alias' % \ + self.__class__.__name__, }, }) + def set_pinned(self, isPinned=True): + return ReturnValue({'BaseResponse': { + 'Ret': -1006, + 'ErrMsg': '%s can not be pinned' % \ + self.__class__.__name__, }, }) + def verify(self): + return ReturnValue({'BaseResponse': { + 'Ret': -1006, + 'ErrMsg': '%s do not need verify' % \ + self.__class__.__name__, }, }) + def get_head_image(self, imageDir=None): + return self.core.get_head_img(self.userName, picDir=imageDir) + def delete_member(self, userName): + return ReturnValue({'BaseResponse': { + 'Ret': -1006, + 'ErrMsg': '%s can not delete member' % \ + self.__class__.__name__, }, }) + def add_member(self, userName): + return ReturnValue({'BaseResponse': { + 'Ret': -1006, + 'ErrMsg': '%s can not add member' % \ + self.__class__.__name__, }, }) + def send_raw_msg(self, msgType, content): + return self.core.send_raw_msg(msgType, content, self.userName) + def send_msg(self, msg='Test Message'): + return self.core.send_msg(msgType, content, self.userName) + def send_file(self, fileDir, mediaId=None): + return self.core.send_file(fileDir, self.userName, mediaId) + def send_image(self, fileDir, mediaId=None): + return self.core.send_image(fileDir, self.userName, mediaId) + def send_video(self, fileDir=None, mediaId=None): + return self.core.send_video(fileDir, self.userName, mediaId) + def send(self, msg, mediaId=None): + return self.core.send(msg, self.userName, mediaId) + def search_member(self, name=None, userName=None, remarkName=None, nickName=None, + wechatAccount=None): + return ReturnValue({'BaseResponse': { + 'Ret': -1006, + 'ErrMsg': '%s do not have members' % \ + self.__class__.__name__, }, }) + def __getattr__(self, value): + value = value[0].upper() + value[1:] + return self.get(value, '') + def __deepcopy__(self, memo): + r = self.__class__([ + (copy.deepcopy(k, memo), copy.deepcopy(v, memo)) + for k, v in self.items()]) + r.core = self.core + return r + def __getstate__(self): + return dict(self) + def __setstate__(self, state): + for k, v in state.items(): + self[k] = v + +class User(AbstractUserDict): + def __init__(self, *args, **kwargs): + super(User, self).__init__(*args, **kwargs) + self.verifyDict = {} + self.memberList = fakeContactList + def update(self): + return self.core.update_friend(self.userName) + def set_alias(self, alias): + return self.core.set_alias(self.userName, alias) + def set_pinned(self, isPinned=True): + return self.core.set_pinned(self.userName, isPinned) + def verify(self): + return self.core.add_friend(**verifyDict) + def __deepcopy__(self, memo): + r = super(User, self).__deepcopy__(memo) + r.verifyDict = copy.deepcopy(self.verifyDict) + return r + + +class MassivePlatform(AbstractUserDict): + def __init__(self, *args, **kwargs): + super(MassivePlatform, self).__init__(*args, **kwargs) + self.memberList = fakeContactList + +class Chatroom(AbstractUserDict): + def __init__(self, *args, **kwargs): + super(Chatroom, self).__init__(*args, **kwargs) + memberList = ContactList() + def init_fn(d): + d.chatroom = self + memberList.set_default_value(init_fn, ChatroomMember) + for rawMember in self.memberList: + memberList.append(rawMember) + self['MemberList'] = memberList + def update(self, detailedMember=False): + return self.core.update_chatroom(self.userName, detailedMember) + def set_alias(self, alias): + return self.core.set_chatroom_name(self.userName, alias) + def set_pinned(self, isPinned=True): + return self.core.set_pinned(self.userName, isPinned) + def delete_member(self, userName): + return self.core.delete_member_from_chatroom(self.userName, userName) + def add_member(self, userName): + return self.core.add_member_into_chatroom(self.userName, userName) + def search_member(self, name=None, userName=None, remarkName=None, nickName=None, + wechatAccount=None): + with self.core.storageClass.updateLock: + if (name or userName or remarkName or nickName or wechatAccount) is None: + return None + elif userName: # return the only userName match + for m in self.memberList: + if m.userName == userName: + return copy.deepcopy(m) + else: + matchDict = { + 'RemarkName' : remarkName, + 'NickName' : nickName, + 'Alias' : wechatAccount, } + for k in ('RemarkName', 'NickName', 'Alias'): + if matchDict[k] is None: + del matchDict[k] + if name: # select based on name + contact = [] + for m in self.memberList: + if any([m.get(k) == name for k in ('RemarkName', 'NickName', 'Alias')]): + contact.append(m) + else: + contact = self.memberList[:] + if matchDict: # select again based on matchDict + friendList = [] + for m in contact: + if all([m.get(k) == v for k, v in matchDict.items()]): + friendList.append(m) + return copy.deepcopy(friendList) + else: + return copy.deepcopy(contact) + +class ChatroomMember(AbstractUserDict): + def __init__(self, *args, **kwargs): + super(AbstractUserDict, self).__init__(*args, **kwargs) + self.core = fakeItchat + self.chatroom = self.fakeChatroom + def get_head_image(self, imageDir=None): + return self.core.get_head_img(self.userName, self.chatroom.userName, picDir=imageDir) + def delete_member(self, userName): + return self.core.delete_member_from_chatroom(self.chatroom.userName, self.userName) + def send_raw_msg(self, msgType, content): + return ReturnValue({'BaseResponse': { + 'Ret': -1006, + 'ErrMsg': '%s can not send message directly' % \ + self.__class__.__name__, }, }) + def send_msg(self, msg='Test Message'): + return ReturnValue({'BaseResponse': { + 'Ret': -1006, + 'ErrMsg': '%s can not send message directly' % \ + self.__class__.__name__, }, }) + def send_file(self, fileDir, mediaId=None): + return ReturnValue({'BaseResponse': { + 'Ret': -1006, + 'ErrMsg': '%s can not send message directly' % \ + self.__class__.__name__, }, }) + def send_image(self, fileDir, mediaId=None): + return ReturnValue({'BaseResponse': { + 'Ret': -1006, + 'ErrMsg': '%s can not send message directly' % \ + self.__class__.__name__, }, }) + def send_video(self, fileDir=None, mediaId=None): + return ReturnValue({'BaseResponse': { + 'Ret': -1006, + 'ErrMsg': '%s can not send message directly' % \ + self.__class__.__name__, }, }) + def send(self, msg, mediaId=None): + return ReturnValue({'BaseResponse': { + 'Ret': -1006, + 'ErrMsg': '%s can not send message directly' % \ + self.__class__.__name__, }, }) + def __deepcopy__(self, memo): + r = super(ChatroomMember, self).__deepcopy__(memo) + r.core = self.core + return r + +ChatroomMember.fakeChatroom = Chatroom() + +def wrap_user_dict(d): + userName = d.get('UserName') + if '@@' in userName: + r = Chatroom(d) + elif d.get('VerifyFlag', 8) & 8 == 0: + r = User(d) + else: + r = MassivePlatform(d) + return r