Add msg['User']

Author: littlecodersh

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: