From 5f729e422710dd216b0ca9ff6cce058df6e897e6 Mon Sep 17 00:00:00 2001
From: Yuxin Wu <ppwwyyxxc@gmail.com>
Date: Sat, 27 Jun 2020 20:04:34 -0700
Subject: [PATCH] support both chatid and nickname (fix #51)

---
 dump-audio.py    |  2 +-
 dump-html.py     |  5 +++--
 wechat/parser.py | 17 +++++++++++++++--
 wechat/res.py    |  4 ++--
 4 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/dump-audio.py b/dump-audio.py
index 47f5fa1..f94be2c 100755
--- a/dump-audio.py
+++ b/dump-audio.py
@@ -24,7 +24,7 @@ if __name__ == '__main__':
     res = Resource(parser, args.res, '')
 
     try:
-        chatid = parser.get_id_by_nickname(args.name)
+        chatid = parser.get_chat_id(args.name)
     except KeyError:
         sys.stderr.write(u"Valid Contacts: {}\n".format('\n'.join(parser.all_chat_nicknames)))
         sys.stderr.write(u"Couldn't find the chat {}.".format(args.name));
diff --git a/dump-html.py b/dump-html.py
index aad680d..3147647 100755
--- a/dump-html.py
+++ b/dump-html.py
@@ -31,15 +31,16 @@ if __name__ == '__main__':
     parser = WeChatDBParser(args.db)
 
     try:
-        chatid = parser.get_id_by_nickname(name)
+        chatid = parser.get_chat_id(args.name)
     except KeyError:
         sys.stderr.write(u"Valid Contacts: {}\n".format(
             u'\n'.join(parser.all_chat_nicknames)))
         sys.stderr.write(u"Couldn't find the chat {}.".format(name));
         sys.exit(1)
+
     res = Resource(parser, args.res, args.avt)
     msgs = parser.msgs_by_chat[chatid]
-    logger.info(f"Number of Messages: {len(msgs)}")
+    logger.info(f"Number of Messages for chatid {chatid}: {len(msgs)}")
     assert len(msgs) > 0
 
     render = HTMLRender(parser, res)
diff --git a/wechat/parser.py b/wechat/parser.py
index 16021d9..3b70cfb 100644
--- a/wechat/parser.py
+++ b/wechat/parser.py
@@ -136,7 +136,7 @@ SELECT {} FROM message
         if values['content']:
             values['content'] = ensure_unicode(values['content'])
         else:
-            values['content'] = u''
+            values['content'] = ''
         values['createTime'] = datetime.fromtimestamp(values['createTime']/ 1000)
         values['chat'] = values['talker']
         try:
@@ -147,7 +147,7 @@ SELECT {} FROM message
                 if values['isSend'] == 1:
                     values['talker'] = self.username
                 elif values['type'] == TYPE_SYSTEM:
-                    values['talker'] = u'SYSTEM'
+                    values['talker'] = 'SYSTEM'
                 else:
                     talker = content[:content.find(':')]
                     values['talker'] = talker
@@ -175,9 +175,22 @@ SELECT {} FROM message
         return [self.contacts[k] for k in self.all_chat_ids if len(self.contacts[k])]
 
     def get_id_by_nickname(self, nickname):
+        """
+        Get chat id by nickname.
+        """
         l = self.contacts_rev[nickname]
         if len(l) == 0:
             raise KeyError("No contacts have nickname {}".format(nickname))
         if len(l) > 1:
             logger.warn("More than one contacts have nickname {}! Using the first contact".format(nickname))
         return l[0]
+
+    def get_chat_id(self, nick_name_or_id):
+        """
+        Get the unique chat id by either chat id itself, or the nickname of the chat.
+        """
+        if nick_name_or_id in self.contacts:
+            return nick_name_or_id
+        else:
+            return self.get_id_by_nickname(nick_name_or_id)
+
diff --git a/wechat/res.py b/wechat/res.py
index e16be77..46f9d80 100644
--- a/wechat/res.py
+++ b/wechat/res.py
@@ -72,8 +72,8 @@ class Resource(object):
     """ multimedia resources in chat"""
     def __init__(self, parser, res_dir, avt_db):
         def check(subdir):
-            assert os.path.isdir(os.path.join(res_dir, subdir)), \
-                    "No such directory: {}".format(subdir)
+            dir_to_check = os.path.join(res_dir, subdir)
+            assert os.path.isdir(dir_to_check), f"No such directory: {dir_to_check}"
         [check(k) for k in ['', IMG_DIRNAME, EMOJI_DIRNAME, VOICE_DIRNAME]]
 
         self.emoji_cache = EmojiCache(