diff --git a/itchat/components/hotreload.py b/itchat/components/hotreload.py index f0a40d04..d34d25ee 100644 --- a/itchat/components/hotreload.py +++ b/itchat/components/hotreload.py @@ -5,6 +5,7 @@ from ..config import VERSION from ..returnvalues import ReturnValue +from ..storage import templates from .contact import update_local_chatrooms, update_local_friends from .messages import produce_msg @@ -50,6 +51,8 @@ def load_login_status(self, fileDir, 'ErrMsg': 'cached status ignored because of version', 'Ret': -1005, }}) self.loginInfo = j['loginInfo'] + self.loginInfo['User'] = templates.User(self.loginInfo['User']) + self.loginInfo['User'].core = self self.s.cookies = requests.utils.cookiejar_from_dict(j['cookies']) self.storageClass.loads(j['storage']) msgList, contactList = self.get_msg() diff --git a/itchat/components/login.py b/itchat/components/login.py index 6b5fb043..7736427c 100644 --- a/itchat/components/login.py +++ b/itchat/components/login.py @@ -248,6 +248,7 @@ def maintain_loop(): else: otherList.append(contact) chatroomMsg = update_local_chatrooms(self, chatroomList) + chatroomMsg['User'] = self.loginInfo['User'] self.msgList.put(chatroomMsg) update_local_friends(self, otherList) retryCount = 0 diff --git a/itchat/components/messages.py b/itchat/components/messages.py index 2c3d0081..dac78c61 100644 --- a/itchat/components/messages.py +++ b/itchat/components/messages.py @@ -8,6 +8,7 @@ from .. import config, utils from ..returnvalues import ReturnValue +from ..storage import templates from .contact import update_local_uin logger = logging.getLogger('itchat') @@ -50,10 +51,29 @@ def produce_msg(core, msgList): rl = [] srl = [40, 43, 50, 52, 53, 9999] for m in msgList: + # get actual opposite + if m['FromUserName'] == core.storageClass.userName: + actualOpposite = m['ToUserName'] + else: + actualOpposite = m['FromUserName'] + # produce basic message if '@@' in m['FromUserName'] or '@@' in m['ToUserName']: produce_group_chat(core, m) else: utils.msg_formatter(m, 'Content') + # set user of msg + if '@@' in actualOpposite: + m['User'] = core.search_chatrooms(userName=actualOpposite) or \ + templates.Chatroom({'UserName': actualOpposite}) + # we don't need to update chatroom here because we have + # updated once when producing basic message + elif actualOpposite in ('filehelper', 'fmessage'): + m['User'] = templates.MassivePlatform({'UserName': actualOpposite}) + else: + m['User'] = core.search_mps(userName=actualOpposite) or \ + core.search_friends(userName=actualOpposite) or \ + templates.User(userName=actualOpposite) + # by default we think there may be a user missing not a mp if m['MsgType'] == 1: # words if m['Url']: regx = r'(.+?\(.+?\))' diff --git a/itchat/components/register.py b/itchat/components/register.py index a07c1111..03add000 100644 --- a/itchat/components/register.py +++ b/itchat/components/register.py @@ -6,6 +6,7 @@ from ..log import set_logging from ..utils import test_connect +from ..storage import templates logger = logging.getLogger('itchat') @@ -46,19 +47,12 @@ def configured_reply(self): except Queue.Empty: pass else: - if msg['FromUserName'] == self.storageClass.userName: - actualOpposite = msg['ToUserName'] - else: - actualOpposite = msg['FromUserName'] - if '@@' in actualOpposite: - replyFn = self.functionDict['GroupChat'].get(msg['Type']) - elif self.search_mps(userName=msg['FromUserName']): - replyFn = self.functionDict['MpChat'].get(msg['Type']) - elif '@' in actualOpposite or \ - actualOpposite in ('filehelper', 'fmessage'): + if isinstance(msg['User'], templates.User): replyFn = self.functionDict['FriendChat'].get(msg['Type']) - else: + elif isinstance(msg['User'], templates.MassivePlatform): replyFn = self.functionDict['MpChat'].get(msg['Type']) + elif isinstance(msg['User'], templates.Chatroom): + replyFn = self.functionDict['GroupChat'].get(msg['Type']) if replyFn is None: r = None else: diff --git a/itchat/storage/__init__.py b/itchat/storage/__init__.py index 22419d07..85a2f791 100644 --- a/itchat/storage/__init__.py +++ b/itchat/storage/__init__.py @@ -40,8 +40,8 @@ def dumps(self): 'chatroomList' : self.chatroomList, 'lastInputUserName' : self.lastInputUserName, } def loads(self, j): - self.userName = j.get('userName', None) - self.nickName = j.get('nickName', None) + 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) @@ -88,19 +88,23 @@ def search_chatrooms(self, name=None, userName=None): with self.updateLock: if userName is not None: for m in self.chatroomList: - if m['UserName'] == userName: return copy.deepcopy(m) + if m['UserName'] == userName: + return copy.deepcopy(m) elif name is not None: matchList = [] for m in self.chatroomList: - if name in m['NickName']: matchList.append(copy.deepcopy(m)) + if name in m['NickName']: + matchList.append(copy.deepcopy(m)) return matchList def search_mps(self, name=None, userName=None): with self.updateLock: if userName is not None: for m in self.mpList: - if m['UserName'] == userName: return copy.deepcopy(m) + if m['UserName'] == userName: + return copy.deepcopy(m) elif name is not None: matchList = [] for m in self.mpList: - if name in m['NickName']: matchList.append(copy.deepcopy(m)) + if name in m['NickName']: + matchList.append(copy.deepcopy(m)) return matchList diff --git a/itchat/storage/templates.py b/itchat/storage/templates.py index e1430e74..524596b2 100644 --- a/itchat/storage/templates.py +++ b/itchat/storage/templates.py @@ -29,7 +29,7 @@ def append(self, value): contact = self.contactClass(value) contact.core = self.core if self.contactInitFn is not None: - contact = self.contactInitFn(contact) + contact = self.contactInitFn(contact) or contact super(ContactList, self).append(contact) def __deepcopy__(self, memo): return self.__class__([copy.deepcopy(v) for v in self]) @@ -38,6 +38,11 @@ def __getstate__(self): def __setstate__(self, state): for v in state: super(ContactList, self).append(pickle.loads(v)) + def __str__(self): + return '[%s]' % ', '.join([repr(v) for v in self]) + def __repr__(self): + return '<%s: %s>' % (self.__class__.__name__.split('.')[-1], + self.__str__()) fakeContactList = ContactList @@ -99,9 +104,9 @@ 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 = 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): @@ -109,6 +114,12 @@ def __getstate__(self): def __setstate__(self, state): for k, v in state.items(): self[k] = v + def __str__(self): + return '{%s}' % ', '.join( + ['%s: %s' % (repr(k),repr(v)) for k,v in self.items()]) + def __repr__(self): + return '<%s: %s>' % (self.__class__.__name__.split('.')[-1], + self.__str__()) class User(AbstractUserDict): def __init__(self, *args, **kwargs): @@ -128,7 +139,6 @@ def __deepcopy__(self, memo): r.verifyDict = copy.deepcopy(self.verifyDict) return r - class MassivePlatform(AbstractUserDict): def __init__(self, *args, **kwargs): super(MassivePlatform, self).__init__(*args, **kwargs) diff --git a/itchat/utils.py b/itchat/utils.py index f5cc7a8e..967ee5d7 100644 --- a/itchat/utils.py +++ b/itchat/utils.py @@ -106,7 +106,8 @@ def search_dict_list(l, key, value): ''' Search a list of dict * return dict with specific value & key ''' for i in l: - if i.get(key) == value: return i + if i.get(key) == value: + return i def print_line(msg, oneLine = False): if oneLine: