import: Move 'process_uploads' and 'process_emojis' to import_util.

zerver/data_import/import_util.py

@@ -2,11 +2,12 @@ import random
 import requests
 import shutil
 import logging
+import os
 
 from typing import List, Dict, Any
 from django.forms.models import model_to_dict
 
-from zerver.models import Realm
+from zerver.models import Realm, RealmEmoji
 from zerver.lib.actions import STREAM_ASSIGNMENT_COLORS as stream_colors
 from zerver.lib.avatar_hash import user_avatar_path_from_ids
 from zerver.lib.parallel import run_parallel
@@ -112,3 +113,89 @@ def process_avatars(avatar_list: List[ZerverFieldsT], avatar_dir: str, realm_id:
 
     logging.info('######### GETTING AVATARS FINISHED #########\n')
     return avatar_list + avatar_original_list
+
+def process_uploads(upload_list: List[ZerverFieldsT], upload_dir: str,
+                    threads: int) -> List[ZerverFieldsT]:
+    """
+    This function downloads the uploads and saves it in the realm's upload directory.
+    Required parameters:
+
+    1. upload_list: List of uploads to be mapped in uploads records.json file
+    2. upload_dir: Folder where the downloaded uploads are saved
+    """
+    def get_uploads(upload: List[str]) -> int:
+        upload_url = upload[0]
+        upload_path = upload[1]
+        upload_path = os.path.join(upload_dir, upload_path)
+
+        response = requests.get(upload_url, stream=True)
+        os.makedirs(os.path.dirname(upload_path), exist_ok=True)
+        with open(upload_path, 'wb') as upload_file:
+            shutil.copyfileobj(response.raw, upload_file)
+        return 0
+
+    logging.info('######### GETTING ATTACHMENTS #########\n')
+    logging.info('DOWNLOADING ATTACHMENTS .......\n')
+    upload_url_list = []
+    for upload in upload_list:
+        upload_url = upload['path']
+        upload_s3_path = upload['s3_path']
+        upload_url_list.append([upload_url, upload_s3_path])
+        upload['path'] = upload_s3_path
+
+    # Run downloads parallely
+    output = []
+    for (status, job) in run_parallel(get_uploads, upload_url_list, threads=threads):
+        output.append(job)
+
+    logging.info('######### GETTING ATTACHMENTS FINISHED #########\n')
+    return upload_list
+
+def process_emojis(zerver_realmemoji: List[ZerverFieldsT], emoji_dir: str,
+                   emoji_url_map: ZerverFieldsT, threads: int) -> List[ZerverFieldsT]:
+    """
+    This function downloads the custom emojis and saves in the output emoji folder.
+    Required parameters:
+
+    1. zerver_realmemoji: List of all RealmEmoji objects to be imported
+    2. emoji_dir: Folder where the downloaded emojis are saved
+    3. emoji_url_map: Maps emoji name to its url
+    """
+    def get_emojis(upload: List[str]) -> int:
+        emoji_url = upload[0]
+        emoji_path = upload[1]
+        upload_emoji_path = os.path.join(emoji_dir, emoji_path)
+
+        response = requests.get(emoji_url, stream=True)
+        os.makedirs(os.path.dirname(upload_emoji_path), exist_ok=True)
+        with open(upload_emoji_path, 'wb') as emoji_file:
+            shutil.copyfileobj(response.raw, emoji_file)
+        return 0
+
+    emoji_records = []
+    upload_emoji_list = []
+    logging.info('######### GETTING EMOJIS #########\n')
+    logging.info('DOWNLOADING EMOJIS .......\n')
+    for emoji in zerver_realmemoji:
+        emoji_url = emoji_url_map[emoji['name']]
+        emoji_path = RealmEmoji.PATH_ID_TEMPLATE.format(
+            realm_id=emoji['realm'],
+            emoji_file_name=emoji['name'])
+
+        upload_emoji_list.append([emoji_url, emoji_path])
+
+        emoji_record = dict(emoji)
+        emoji_record['path'] = emoji_path
+        emoji_record['s3_path'] = emoji_path
+        emoji_record['realm_id'] = emoji_record['realm']
+        emoji_record.pop('realm')
+
+        emoji_records.append(emoji_record)
+
+    # Run downloads parallely
+    output = []
+    for (status, job) in run_parallel(get_emojis, upload_emoji_list, threads=threads):
+        output.append(job)
+
+    logging.info('######### GETTING EMOJIS FINISHED #########\n')
+    return emoji_records

zerver/data_import/slack.py

@@ -21,7 +21,8 @@ from zerver.models import Reaction, RealmEmoji, Realm, UserProfile, Recipient
 from zerver.data_import.slack_message_conversion import convert_to_zulip_markdown, \
     get_user_full_name
 from zerver.data_import.import_util import ZerverFieldsT, build_zerver_realm, \
-    build_avatar, build_subscription, build_recipient, process_avatars
+    build_avatar, build_subscription, build_recipient, process_avatars, \
+    process_uploads, process_emojis
 from zerver.lib.parallel import run_parallel
 from zerver.lib.upload import random_name, sanitize_name
 from zerver.lib.export import MESSAGE_BATCH_CHUNK_SIZE
@@ -843,83 +844,6 @@ def do_convert_data(slack_zip_file: str, output_dir: str, token: str, threads: i
     logging.info('######### DATA CONVERSION FINISHED #########\n')
     logging.info("Zulip data dump created at %s" % (output_dir))
 
-def process_emojis(zerver_realmemoji: List[ZerverFieldsT], emoji_dir: str,
-                   emoji_url_map: ZerverFieldsT, threads: int) -> List[ZerverFieldsT]:
-    """
-    This function gets the custom emojis and saves in the output emoji folder
-    """
-    def get_emojis(upload: List[str]) -> int:
-        slack_emoji_url = upload[0]
-        emoji_path = upload[1]
-        upload_emoji_path = os.path.join(emoji_dir, emoji_path)
-
-        response = requests.get(slack_emoji_url, stream=True)
-        os.makedirs(os.path.dirname(upload_emoji_path), exist_ok=True)
-        with open(upload_emoji_path, 'wb') as emoji_file:
-            shutil.copyfileobj(response.raw, emoji_file)
-        return 0
-
-    emoji_records = []
-    upload_emoji_list = []
-    logging.info('######### GETTING EMOJIS #########\n')
-    logging.info('DOWNLOADING EMOJIS .......\n')
-    for emoji in zerver_realmemoji:
-        slack_emoji_url = emoji_url_map[emoji['name']]
-        emoji_path = RealmEmoji.PATH_ID_TEMPLATE.format(
-            realm_id=emoji['realm'],
-            emoji_file_name=emoji['name'])
-
-        upload_emoji_list.append([slack_emoji_url, emoji_path])
-
-        emoji_record = dict(emoji)
-        emoji_record['path'] = emoji_path
-        emoji_record['s3_path'] = emoji_path
-        emoji_record['realm_id'] = emoji_record['realm']
-        emoji_record.pop('realm')
-
-        emoji_records.append(emoji_record)
-
-    # Run downloads parallely
-    output = []
-    for (status, job) in run_parallel(get_emojis, upload_emoji_list, threads=threads):
-        output.append(job)
-
-    logging.info('######### GETTING EMOJIS FINISHED #########\n')
-    return emoji_records
-
-def process_uploads(upload_list: List[ZerverFieldsT], upload_dir: str,
-                    threads: int) -> List[ZerverFieldsT]:
-    """
-    This function gets the uploads and saves it in the realm's upload directory
-    """
-    def get_uploads(upload: List[str]) -> int:
-        upload_url = upload[0]
-        upload_path = upload[1]
-        upload_path = os.path.join(upload_dir, upload_path)
-
-        response = requests.get(upload_url, stream=True)
-        os.makedirs(os.path.dirname(upload_path), exist_ok=True)
-        with open(upload_path, 'wb') as upload_file:
-            shutil.copyfileobj(response.raw, upload_file)
-        return 0
-
-    logging.info('######### GETTING ATTACHMENTS #########\n')
-    logging.info('DOWNLOADING ATTACHMENTS .......\n')
-    upload_url_list = []
-    for upload in upload_list:
-        upload_url = upload['path']
-        upload_s3_path = upload['s3_path']
-        upload_url_list.append([upload_url, upload_s3_path])
-        upload['path'] = upload_s3_path
-
-    # Run downloads parallely
-    output = []
-    for (status, job) in run_parallel(get_uploads, upload_url_list, threads=threads):
-        output.append(job)
-
-    logging.info('######### GETTING ATTACHMENTS FINISHED #########\n')
-    return upload_list
-
 def get_data_file(path: str) -> Any:
     data = json.load(open(path))
     return data