Add msg['User']

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()

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

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'(.+?\(.+?\))'

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:

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

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,16 +104,22 @@ 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):
         return dict(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)

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: