From 032af6cdcf3377e2736badaaaee36fdc91673be5 Mon Sep 17 00:00:00 2001
From: KnugiHK <24708955+KnugiHK@users.noreply.github.com>
Date: Wed, 21 Jun 2023 17:01:14 +0800
Subject: [PATCH] Refactor

---
 Whatsapp_Chat_Exporter/data_model.py |  2 +-
 Whatsapp_Chat_Exporter/extract.py    | 31 ++++----------------------
 Whatsapp_Chat_Exporter/utility.py    | 33 ++++++++++++++++++++++++++++
 3 files changed, 38 insertions(+), 28 deletions(-)

diff --git a/Whatsapp_Chat_Exporter/data_model.py b/Whatsapp_Chat_Exporter/data_model.py
index 6c912a5..07c95cc 100644
--- a/Whatsapp_Chat_Exporter/data_model.py
+++ b/Whatsapp_Chat_Exporter/data_model.py
@@ -1,7 +1,6 @@
 import os
 from datetime import datetime
 from typing import Union
-from Whatsapp_Chat_Exporter.utility import Device
 
 
 class ChatStore():
@@ -12,6 +11,7 @@ class ChatStore():
         self.messages = {}
         self.type = type
         if media is not None:
+            from Whatsapp_Chat_Exporter.utility import Device
             if self.type == Device.IOS:
                 self.my_avatar = os.path.join(media, "Media/Profile/Photo.jpg")
             elif self.type == Device.ANDROID:
diff --git a/Whatsapp_Chat_Exporter/extract.py b/Whatsapp_Chat_Exporter/extract.py
index 038579e..b88b49e 100644
--- a/Whatsapp_Chat_Exporter/extract.py
+++ b/Whatsapp_Chat_Exporter/extract.py
@@ -13,8 +13,8 @@ from mimetypes import MimeTypes
 from hashlib import sha256
 from base64 import b64decode, b64encode
 from Whatsapp_Chat_Exporter.data_model import ChatStore, Message
-from Whatsapp_Chat_Exporter.utility import MAX_SIZE, ROW_SIZE, Device, determine_metadata
-from Whatsapp_Chat_Exporter.utility import rendering, sanitize_except, determine_day, Crypt, get_file_name
+from Whatsapp_Chat_Exporter.utility import MAX_SIZE, ROW_SIZE, determine_metadata, get_status_location
+from Whatsapp_Chat_Exporter.utility import rendering, Crypt, Device, get_file_name, setup_template
 from Whatsapp_Chat_Exporter.utility import brute_force_offset, CRYPT14_OFFSETS
 
 try:
@@ -639,20 +639,7 @@ def create_html(
         maximum_size=None,
         no_avatar=False
     ):
-    if template is None:
-        template_dir = os.path.dirname(__file__)
-        template_file = "whatsapp.html"
-    else:
-        template_dir = os.path.dirname(template)
-        template_file = os.path.basename(template)
-    template_loader = jinja2.FileSystemLoader(searchpath=template_dir)
-    template_env = jinja2.Environment(loader=template_loader, autoescape=True)
-    template_env.globals.update(
-        determine_day=determine_day,
-        no_avatar=no_avatar
-    )
-    template_env.filters['sanitize_except'] = sanitize_except
-    template = template_env.get_template(template_file)
+    template = setup_template(template, no_avatar)
 
     total_row_number = len(data)
     print(f"\nGenerating chats...(0/{total_row_number})", end="\r")
@@ -660,17 +647,7 @@ def create_html(
     if not os.path.isdir(output_folder):
         os.mkdir(output_folder)
 
-    w3css = "https://www.w3schools.com/w3css/4/w3.css"
-    if offline_static:
-        import urllib.request
-        static_folder = os.path.join(output_folder, offline_static)
-        if not os.path.isdir(static_folder):
-            os.mkdir(static_folder)
-        w3css_path = os.path.join(static_folder, "w3.css")
-        if not os.path.isfile(w3css_path):
-            with urllib.request.urlopen(w3css) as resp:
-                with open(w3css_path, "wb") as f: f.write(resp.read())
-        w3css = os.path.join(offline_static, "w3.css")
+    w3css = get_status_location(output_folder, offline_static)
 
     for current, contact in enumerate(data):
         chat = data[contact]
diff --git a/Whatsapp_Chat_Exporter/utility.py b/Whatsapp_Chat_Exporter/utility.py
index 3128336..bb479d0 100644
--- a/Whatsapp_Chat_Exporter/utility.py
+++ b/Whatsapp_Chat_Exporter/utility.py
@@ -1,4 +1,6 @@
+import jinja2
 import json
+import os
 from bleach import clean as sanitize
 from markupsafe import Markup
 from datetime import datetime
@@ -254,5 +256,36 @@ def determine_metadata(content, init_msg):
     return msg
 
 
+def get_status_location(output_folder, offline_static):
+    w3css = "https://www.w3schools.com/w3css/4/w3.css"
+    if not offline_static:
+        return w3css
+    import urllib.request
+    static_folder = os.path.join(output_folder, offline_static)
+    if not os.path.isdir(static_folder):
+        os.mkdir(static_folder)
+    w3css_path = os.path.join(static_folder, "w3.css")
+    if not os.path.isfile(w3css_path):
+        with urllib.request.urlopen(w3css) as resp:
+            with open(w3css_path, "wb") as f: f.write(resp.read())
+    w3css = os.path.join(offline_static, "w3.css")
+
+
+def setup_template(template, no_avatar):
+    if template is None:
+        template_dir = os.path.dirname(__file__)
+        template_file = "whatsapp.html"
+    else:
+        template_dir = os.path.dirname(template)
+        template_file = os.path.basename(template)
+    template_loader = jinja2.FileSystemLoader(searchpath=template_dir)
+    template_env = jinja2.Environment(loader=template_loader, autoescape=True)
+    template_env.globals.update(
+        determine_day=determine_day,
+        no_avatar=no_avatar
+    )
+    template_env.filters['sanitize_except'] = sanitize_except
+    return template_env.get_template(template_file)
+
 # iOS Specific
 APPLE_TIME = datetime.timestamp(datetime(2001, 1, 1))