__init__.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117
  1. import os, time, copy
  2. from threading import Lock
  3. from .messagequeue import Queue
  4. from .templates import (
  5. ContactList, AbstractUserDict, User,
  6. MassivePlatform, Chatroom, ChatroomMember)
  7. def contact_change(fn):
  8. def _contact_change(core, *args, **kwargs):
  9. with core.storageClass.updateLock:
  10. return fn(core, *args, **kwargs)
  11. return _contact_change
  12. class Storage(object):
  13. def __init__(self, core):
  14. self.userName = None
  15. self.nickName = None
  16. self.updateLock = Lock()
  17. self.memberList = ContactList()
  18. self.mpList = ContactList()
  19. self.chatroomList = ContactList()
  20. self.msgList = Queue(-1)
  21. self.lastInputUserName = None
  22. self.memberList.set_default_value(contactClass=User)
  23. self.memberList.core = core
  24. self.mpList.set_default_value(contactClass=MassivePlatform)
  25. self.mpList.core = core
  26. self.chatroomList.set_default_value(contactClass=Chatroom)
  27. self.chatroomList.core = core
  28. def dumps(self):
  29. return {
  30. 'userName' : self.userName,
  31. 'nickName' : self.nickName,
  32. 'memberList' : self.memberList,
  33. 'mpList' : self.mpList,
  34. 'chatroomList' : self.chatroomList,
  35. 'lastInputUserName' : self.lastInputUserName, }
  36. def loads(self, j):
  37. self.userName = j.get('userName', None)
  38. self.nickName = j.get('nickName', None)
  39. del self.memberList[:]
  40. for i in j.get('memberList', []):
  41. self.memberList.append(i)
  42. del self.mpList[:]
  43. for i in j.get('mpList', []):
  44. self.mpList.append(i)
  45. del self.chatroomList[:]
  46. for i in j.get('chatroomList', []):
  47. self.chatroomList.append(i)
  48. # I tried to solve everything in pickle
  49. # but this way is easier and more storage-saving
  50. for chatroom in self.chatroomList:
  51. if 'MemberList' in chatroom:
  52. for member in chatroom['MemberList']:
  53. member.core = chatroom.core
  54. member.chatroom = chatroom
  55. if 'Self' in chatroom:
  56. chatroom['Self'].core = chatroom.core
  57. chatroom['Self'].chatroom = chatroom
  58. self.lastInputUserName = j.get('lastInputUserName', None)
  59. def search_friends(self, name=None, userName=None, remarkName=None, nickName=None,
  60. wechatAccount=None):
  61. with self.updateLock:
  62. if (name or userName or remarkName or nickName or wechatAccount) is None:
  63. return copy.deepcopy(self.memberList[0]) # my own account
  64. elif userName: # return the only userName match
  65. for m in self.memberList:
  66. if m['UserName'] == userName:
  67. return copy.deepcopy(m)
  68. else:
  69. matchDict = {
  70. 'RemarkName' : remarkName,
  71. 'NickName' : nickName,
  72. 'Alias' : wechatAccount, }
  73. for k in ('RemarkName', 'NickName', 'Alias'):
  74. if matchDict[k] is None:
  75. del matchDict[k]
  76. if name: # select based on name
  77. contact = []
  78. for m in self.memberList:
  79. if any([m.get(k) == name for k in ('RemarkName', 'NickName', 'Alias')]):
  80. contact.append(m)
  81. else:
  82. contact = self.memberList[:]
  83. if matchDict: # select again based on matchDict
  84. friendList = []
  85. for m in contact:
  86. if all([m.get(k) == v for k, v in matchDict.items()]):
  87. friendList.append(m)
  88. return copy.deepcopy(friendList)
  89. else:
  90. return copy.deepcopy(contact)
  91. def search_chatrooms(self, name=None, userName=None):
  92. with self.updateLock:
  93. if userName is not None:
  94. for m in self.chatroomList:
  95. if m['UserName'] == userName:
  96. return copy.deepcopy(m)
  97. elif name is not None:
  98. matchList = []
  99. for m in self.chatroomList:
  100. if name in m['NickName']:
  101. matchList.append(copy.deepcopy(m))
  102. return matchList
  103. def search_mps(self, name=None, userName=None):
  104. with self.updateLock:
  105. if userName is not None:
  106. for m in self.mpList:
  107. if m['UserName'] == userName:
  108. return copy.deepcopy(m)
  109. elif name is not None:
  110. matchList = []
  111. for m in self.mpList:
  112. if name in m['NickName']:
  113. matchList.append(copy.deepcopy(m))
  114. return matchList