python: Reformat with Black, except quotes.

Signed-off-by: Anders Kaseorg <anders@zulip.com>
This commit is contained in:
Anders Kaseorg
2021-02-11 23:19:30 -08:00
committed by Tim Abbott
parent 5028c081cb
commit 11741543da
817 changed files with 44952 additions and 24860 deletions

View File

@@ -28,16 +28,18 @@ from zerver.models import (
# stubs
ZerverFieldsT = Dict[str, Any]
class SubscriberHandler:
def __init__(self) -> None:
self.stream_info: Dict[int, Set[int]] = {}
self.huddle_info: Dict[int, Set[int]] = {}
def set_info(self,
users: Set[int],
stream_id: Optional[int]=None,
huddle_id: Optional[int]=None,
) -> None:
def set_info(
self,
users: Set[int],
stream_id: Optional[int] = None,
huddle_id: Optional[int] = None,
) -> None:
if stream_id is not None:
self.stream_info[stream_id] = users
elif huddle_id is not None:
@@ -45,9 +47,9 @@ class SubscriberHandler:
else:
raise AssertionError("stream_id or huddle_id is required")
def get_users(self,
stream_id: Optional[int]=None,
huddle_id: Optional[int]=None) -> Set[int]:
def get_users(
self, stream_id: Optional[int] = None, huddle_id: Optional[int] = None
) -> Set[int]:
if stream_id is not None:
return self.stream_info[stream_id]
elif huddle_id is not None:
@@ -55,28 +57,37 @@ class SubscriberHandler:
else:
raise AssertionError("stream_id or huddle_id is required")
def build_zerver_realm(realm_id: int, realm_subdomain: str, time: float,
other_product: str) -> List[ZerverFieldsT]:
realm = Realm(id=realm_id, date_created=time,
name=realm_subdomain, string_id=realm_subdomain,
description=f"Organization imported from {other_product}!")
def build_zerver_realm(
realm_id: int, realm_subdomain: str, time: float, other_product: str
) -> List[ZerverFieldsT]:
realm = Realm(
id=realm_id,
date_created=time,
name=realm_subdomain,
string_id=realm_subdomain,
description=f"Organization imported from {other_product}!",
)
auth_methods = [[flag[0], flag[1]] for flag in realm.authentication_methods]
realm_dict = model_to_dict(realm, exclude='authentication_methods')
realm_dict['authentication_methods'] = auth_methods
return[realm_dict]
return [realm_dict]
def build_user_profile(avatar_source: str,
date_joined: Any,
delivery_email: str,
email: str,
full_name: str,
id: int,
is_active: bool,
role: int,
is_mirror_dummy: bool,
realm_id: int,
short_name: str,
timezone: Optional[str]) -> ZerverFieldsT:
def build_user_profile(
avatar_source: str,
date_joined: Any,
delivery_email: str,
email: str,
full_name: str,
id: int,
is_active: bool,
role: int,
is_mirror_dummy: bool,
realm_id: int,
short_name: str,
timezone: Optional[str],
) -> ZerverFieldsT:
obj = UserProfile(
avatar_source=avatar_source,
date_joined=date_joined,
@@ -100,8 +111,15 @@ def build_user_profile(avatar_source: str,
dct['short_name'] = short_name
return dct
def build_avatar(zulip_user_id: int, realm_id: int, email: str, avatar_url: str,
timestamp: Any, avatar_list: List[ZerverFieldsT]) -> None:
def build_avatar(
zulip_user_id: int,
realm_id: int,
email: str,
avatar_url: str,
timestamp: Any,
avatar_list: List[ZerverFieldsT],
) -> None:
avatar = dict(
path=avatar_url, # Save original avatar URL here, which is downloaded later
realm_id=realm_id,
@@ -110,14 +128,16 @@ def build_avatar(zulip_user_id: int, realm_id: int, email: str, avatar_url: str,
last_modified=timestamp,
user_profile_email=email,
s3_path="",
size="")
size="",
)
avatar_list.append(avatar)
def make_subscriber_map(zerver_subscription: List[ZerverFieldsT]) -> Dict[int, Set[int]]:
'''
"""
This can be convenient for building up UserMessage
rows.
'''
"""
subscriber_map: Dict[int, Set[int]] = {}
for sub in zerver_subscription:
user_id = sub['user_profile']
@@ -128,10 +148,13 @@ def make_subscriber_map(zerver_subscription: List[ZerverFieldsT]) -> Dict[int, S
return subscriber_map
def make_user_messages(zerver_message: List[ZerverFieldsT],
subscriber_map: Dict[int, Set[int]],
is_pm_data: bool,
mention_map: Dict[int, Set[int]]) -> List[ZerverFieldsT]:
def make_user_messages(
zerver_message: List[ZerverFieldsT],
subscriber_map: Dict[int, Set[int]],
is_pm_data: bool,
mention_map: Dict[int, Set[int]],
) -> List[ZerverFieldsT]:
zerver_usermessage = []
@@ -155,46 +178,38 @@ def make_user_messages(zerver_message: List[ZerverFieldsT],
return zerver_usermessage
def build_subscription(recipient_id: int, user_id: int,
subscription_id: int) -> ZerverFieldsT:
subscription = Subscription(
color=random.choice(stream_colors),
id=subscription_id)
def build_subscription(recipient_id: int, user_id: int, subscription_id: int) -> ZerverFieldsT:
subscription = Subscription(color=random.choice(stream_colors), id=subscription_id)
subscription_dict = model_to_dict(subscription, exclude=['user_profile', 'recipient_id'])
subscription_dict['user_profile'] = user_id
subscription_dict['recipient'] = recipient_id
return subscription_dict
def build_public_stream_subscriptions(
zerver_userprofile: List[ZerverFieldsT],
zerver_recipient: List[ZerverFieldsT],
zerver_stream: List[ZerverFieldsT]) -> List[ZerverFieldsT]:
'''
zerver_userprofile: List[ZerverFieldsT],
zerver_recipient: List[ZerverFieldsT],
zerver_stream: List[ZerverFieldsT],
) -> List[ZerverFieldsT]:
"""
This function was only used for HipChat, but it may apply to
future conversions. We often did't get full subscriber data in
the HipChat export, so this function just autosubscribes all
users to every public stream. This returns a list of Subscription
dicts.
'''
"""
subscriptions: List[ZerverFieldsT] = []
public_stream_ids = {
stream['id']
for stream in zerver_stream
if not stream['invite_only']
}
public_stream_ids = {stream['id'] for stream in zerver_stream if not stream['invite_only']}
public_stream_recipient_ids = {
recipient['id']
for recipient in zerver_recipient
if recipient['type'] == Recipient.STREAM
and recipient['type_id'] in public_stream_ids
if recipient['type'] == Recipient.STREAM and recipient['type_id'] in public_stream_ids
}
user_ids = [
user['id']
for user in zerver_userprofile
]
user_ids = [user['id'] for user in zerver_userprofile]
for recipient_id in public_stream_recipient_ids:
for user_id in user_ids:
@@ -207,10 +222,12 @@ def build_public_stream_subscriptions(
return subscriptions
def build_stream_subscriptions(
get_users: Callable[..., Set[int]],
zerver_recipient: List[ZerverFieldsT],
zerver_stream: List[ZerverFieldsT]) -> List[ZerverFieldsT]:
get_users: Callable[..., Set[int]],
zerver_recipient: List[ZerverFieldsT],
zerver_stream: List[ZerverFieldsT],
) -> List[ZerverFieldsT]:
subscriptions: List[ZerverFieldsT] = []
@@ -219,8 +236,7 @@ def build_stream_subscriptions(
recipient_map = {
recipient['id']: recipient['type_id'] # recipient_id -> stream_id
for recipient in zerver_recipient
if recipient['type'] == Recipient.STREAM
and recipient['type_id'] in stream_ids
if recipient['type'] == Recipient.STREAM and recipient['type_id'] in stream_ids
}
for recipient_id, stream_id in recipient_map.items():
@@ -235,10 +251,12 @@ def build_stream_subscriptions(
return subscriptions
def build_huddle_subscriptions(
get_users: Callable[..., Set[int]],
zerver_recipient: List[ZerverFieldsT],
zerver_huddle: List[ZerverFieldsT]) -> List[ZerverFieldsT]:
get_users: Callable[..., Set[int]],
zerver_recipient: List[ZerverFieldsT],
zerver_huddle: List[ZerverFieldsT],
) -> List[ZerverFieldsT]:
subscriptions: List[ZerverFieldsT] = []
@@ -247,8 +265,7 @@ def build_huddle_subscriptions(
recipient_map = {
recipient['id']: recipient['type_id'] # recipient_id -> stream_id
for recipient in zerver_recipient
if recipient['type'] == Recipient.HUDDLE
and recipient['type_id'] in huddle_ids
if recipient['type'] == Recipient.HUDDLE and recipient['type_id'] in huddle_ids
}
for recipient_id, huddle_id in recipient_map.items():
@@ -263,14 +280,13 @@ def build_huddle_subscriptions(
return subscriptions
def build_personal_subscriptions(zerver_recipient: List[ZerverFieldsT]) -> List[ZerverFieldsT]:
subscriptions: List[ZerverFieldsT] = []
personal_recipients = [
recipient
for recipient in zerver_recipient
if recipient['type'] == Recipient.PERSONAL
recipient for recipient in zerver_recipient if recipient['type'] == Recipient.PERSONAL
]
for recipient in personal_recipients:
@@ -285,6 +301,7 @@ def build_personal_subscriptions(zerver_recipient: List[ZerverFieldsT]) -> List[
return subscriptions
def build_recipient(type_id: int, recipient_id: int, type: int) -> ZerverFieldsT:
recipient = Recipient(
type_id=type_id, # stream id
@@ -294,14 +311,17 @@ def build_recipient(type_id: int, recipient_id: int, type: int) -> ZerverFieldsT
recipient_dict = model_to_dict(recipient)
return recipient_dict
def build_recipients(zerver_userprofile: Iterable[ZerverFieldsT],
zerver_stream: Iterable[ZerverFieldsT],
zerver_huddle: Iterable[ZerverFieldsT] = []) -> List[ZerverFieldsT]:
'''
def build_recipients(
zerver_userprofile: Iterable[ZerverFieldsT],
zerver_stream: Iterable[ZerverFieldsT],
zerver_huddle: Iterable[ZerverFieldsT] = [],
) -> List[ZerverFieldsT]:
"""
This function was only used HipChat import, this function may be
required for future conversions. The Slack and Gitter conversions do it more
tightly integrated with creating other objects.
'''
"""
recipients = []
@@ -339,36 +359,44 @@ def build_recipients(zerver_userprofile: Iterable[ZerverFieldsT],
recipients.append(recipient_dict)
return recipients
def build_realm(zerver_realm: List[ZerverFieldsT], realm_id: int,
domain_name: str) -> ZerverFieldsT:
realm = dict(zerver_client=[{"name": "populate_db", "id": 1},
{"name": "website", "id": 2},
{"name": "API", "id": 3}],
zerver_customprofilefield=[],
zerver_customprofilefieldvalue=[],
zerver_userpresence=[], # shows last logged in data, which is not available
zerver_userprofile_mirrordummy=[],
zerver_realmdomain=[{"realm": realm_id,
"allow_subdomains": False,
"domain": domain_name,
"id": realm_id}],
zerver_useractivity=[],
zerver_realm=zerver_realm,
zerver_huddle=[],
zerver_userprofile_crossrealm=[],
zerver_useractivityinterval=[],
zerver_reaction=[],
zerver_realmemoji=[],
zerver_realmfilter=[])
def build_realm(
zerver_realm: List[ZerverFieldsT], realm_id: int, domain_name: str
) -> ZerverFieldsT:
realm = dict(
zerver_client=[
{"name": "populate_db", "id": 1},
{"name": "website", "id": 2},
{"name": "API", "id": 3},
],
zerver_customprofilefield=[],
zerver_customprofilefieldvalue=[],
zerver_userpresence=[], # shows last logged in data, which is not available
zerver_userprofile_mirrordummy=[],
zerver_realmdomain=[
{"realm": realm_id, "allow_subdomains": False, "domain": domain_name, "id": realm_id}
],
zerver_useractivity=[],
zerver_realm=zerver_realm,
zerver_huddle=[],
zerver_userprofile_crossrealm=[],
zerver_useractivityinterval=[],
zerver_reaction=[],
zerver_realmemoji=[],
zerver_realmfilter=[],
)
return realm
def build_usermessages(zerver_usermessage: List[ZerverFieldsT],
subscriber_map: Dict[int, Set[int]],
recipient_id: int,
mentioned_user_ids: List[int],
message_id: int,
is_private: bool,
long_term_idle: AbstractSet[int] = set()) -> Tuple[int, int]:
def build_usermessages(
zerver_usermessage: List[ZerverFieldsT],
subscriber_map: Dict[int, Set[int]],
recipient_id: int,
mentioned_user_ids: List[int],
message_id: int,
is_private: bool,
long_term_idle: AbstractSet[int] = set(),
) -> Tuple[int, int]:
user_ids = subscriber_map.get(recipient_id, set())
user_messages_created = 0
@@ -393,10 +421,10 @@ def build_usermessages(zerver_usermessage: List[ZerverFieldsT],
zerver_usermessage.append(usermessage)
return (user_messages_created, user_messages_skipped)
def build_user_message(user_id: int,
message_id: int,
is_private: bool,
is_mentioned: bool) -> ZerverFieldsT:
def build_user_message(
user_id: int, message_id: int, is_private: bool, is_mentioned: bool
) -> ZerverFieldsT:
flags_mask = 1 # For read
if is_mentioned:
flags_mask += 8 # For mentioned
@@ -413,17 +441,21 @@ def build_user_message(user_id: int,
)
return usermessage
def build_defaultstream(realm_id: int, stream_id: int,
defaultstream_id: int) -> ZerverFieldsT:
defaultstream = dict(
stream=stream_id,
realm=realm_id,
id=defaultstream_id)
def build_defaultstream(realm_id: int, stream_id: int, defaultstream_id: int) -> ZerverFieldsT:
defaultstream = dict(stream=stream_id, realm=realm_id, id=defaultstream_id)
return defaultstream
def build_stream(date_created: Any, realm_id: int, name: str,
description: str, stream_id: int, deactivated: bool=False,
invite_only: bool=False) -> ZerverFieldsT:
def build_stream(
date_created: Any,
realm_id: int,
name: str,
description: str,
stream_id: int,
deactivated: bool = False,
invite_only: bool = False,
) -> ZerverFieldsT:
stream = Stream(
name=name,
deactivated=deactivated,
@@ -431,22 +463,32 @@ def build_stream(date_created: Any, realm_id: int, name: str,
# We don't set rendered_description here; it'll be added on import
date_created=date_created,
invite_only=invite_only,
id=stream_id)
stream_dict = model_to_dict(stream,
exclude=['realm'])
id=stream_id,
)
stream_dict = model_to_dict(stream, exclude=['realm'])
stream_dict['realm'] = realm_id
return stream_dict
def build_huddle(huddle_id: int) -> ZerverFieldsT:
huddle = Huddle(
id=huddle_id,
)
return model_to_dict(huddle)
def build_message(topic_name: str, date_sent: float, message_id: int, content: str,
rendered_content: Optional[str], user_id: int, recipient_id: int,
has_image: bool=False, has_link: bool=False,
has_attachment: bool=True) -> ZerverFieldsT:
def build_message(
topic_name: str,
date_sent: float,
message_id: int,
content: str,
rendered_content: Optional[str],
user_id: int,
recipient_id: int,
has_image: bool = False,
has_link: bool = False,
has_attachment: bool = True,
) -> ZerverFieldsT:
zulip_message = Message(
rendered_content_version=1, # this is Zulip specific
date_sent=date_sent,
@@ -455,19 +497,27 @@ def build_message(topic_name: str, date_sent: float, message_id: int, content: s
rendered_content=rendered_content,
has_image=has_image,
has_attachment=has_attachment,
has_link=has_link)
has_link=has_link,
)
zulip_message.set_topic_name(topic_name)
zulip_message_dict = model_to_dict(zulip_message,
exclude=['recipient', 'sender', 'sending_client'])
zulip_message_dict = model_to_dict(
zulip_message, exclude=['recipient', 'sender', 'sending_client']
)
zulip_message_dict['sender'] = user_id
zulip_message_dict['sending_client'] = 1
zulip_message_dict['recipient'] = recipient_id
return zulip_message_dict
def build_attachment(realm_id: int, message_ids: Set[int],
user_id: int, fileinfo: ZerverFieldsT, s3_path: str,
zerver_attachment: List[ZerverFieldsT]) -> None:
def build_attachment(
realm_id: int,
message_ids: Set[int],
user_id: int,
fileinfo: ZerverFieldsT,
s3_path: str,
zerver_attachment: List[ZerverFieldsT],
) -> None:
"""
This function should be passed a 'fileinfo' dictionary, which contains
information about 'size', 'created' (created time) and ['name'] (filename).
@@ -480,16 +530,17 @@ def build_attachment(realm_id: int, message_ids: Set[int],
create_time=fileinfo['created'],
is_realm_public=True,
path_id=s3_path,
file_name=fileinfo['name'])
file_name=fileinfo['name'],
)
attachment_dict = model_to_dict(attachment,
exclude=['owner', 'messages', 'realm'])
attachment_dict = model_to_dict(attachment, exclude=['owner', 'messages', 'realm'])
attachment_dict['owner'] = user_id
attachment_dict['messages'] = list(message_ids)
attachment_dict['realm'] = realm_id
zerver_attachment.append(attachment_dict)
def get_avatar(avatar_dir: str, size_url_suffix: str, avatar_upload_item: List[str]) -> None:
avatar_url = avatar_upload_item[0]
@@ -501,8 +552,14 @@ def get_avatar(avatar_dir: str, size_url_suffix: str, avatar_upload_item: List[s
shutil.copyfileobj(response.raw, image_file)
shutil.copy(image_path, original_image_path)
def process_avatars(avatar_list: List[ZerverFieldsT], avatar_dir: str, realm_id: int,
threads: int, size_url_suffix: str='') -> List[ZerverFieldsT]:
def process_avatars(
avatar_list: List[ZerverFieldsT],
avatar_dir: str,
realm_id: int,
threads: int,
size_url_suffix: str = '',
) -> List[ZerverFieldsT]:
"""
This function gets the avatar of the user and saves it in the
user's avatar directory with both the extensions '.png' and '.original'
@@ -540,21 +597,21 @@ def process_avatars(avatar_list: List[ZerverFieldsT], avatar_dir: str, realm_id:
avatar_original_list.append(avatar_original)
# Run downloads in parallel
run_parallel_wrapper(partial(get_avatar, avatar_dir, size_url_suffix), avatar_upload_list, threads=threads)
run_parallel_wrapper(
partial(get_avatar, avatar_dir, size_url_suffix), avatar_upload_list, threads=threads
)
logging.info('######### GETTING AVATARS FINISHED #########\n')
return avatar_list + avatar_original_list
def write_avatar_png(avatar_folder: str,
realm_id: int,
user_id: int,
bits: bytes) -> ZerverFieldsT:
'''
def write_avatar_png(avatar_folder: str, realm_id: int, user_id: int, bits: bytes) -> ZerverFieldsT:
"""
Use this function for conversions like HipChat where
the bits for the .png file come in something like
a users.json file, and where we don't have to
fetch avatar images externally.
'''
"""
avatar_hash = user_avatar_path_from_ids(
user_profile_id=user_id,
realm_id=realm_id,
@@ -578,15 +635,20 @@ def write_avatar_png(avatar_folder: str,
return metadata
ListJobData = TypeVar('ListJobData')
def wrapping_function(f: Callable[[ListJobData], None], item: ListJobData) -> None:
try:
f(item)
except Exception:
logging.exception("Error processing item: %s", item, stack_info=True)
def run_parallel_wrapper(f: Callable[[ListJobData], None], full_items: List[ListJobData],
threads: int=6) -> None:
def run_parallel_wrapper(
f: Callable[[ListJobData], None], full_items: List[ListJobData], threads: int = 6
) -> None:
logging.info("Distributing %s items across %s threads", len(full_items), threads)
with multiprocessing.Pool(threads) as p:
@@ -596,6 +658,7 @@ def run_parallel_wrapper(f: Callable[[ListJobData], None], full_items: List[List
if count % 1000 == 0:
logging.info("Finished %s items", count)
def get_uploads(upload_dir: str, upload: List[str]) -> None:
upload_url = upload[0]
upload_path = upload[1]
@@ -606,8 +669,10 @@ def get_uploads(upload_dir: str, upload: List[str]) -> None:
with open(upload_path, 'wb') as upload_file:
shutil.copyfileobj(response.raw, upload_file)
def process_uploads(upload_list: List[ZerverFieldsT], upload_dir: str,
threads: int) -> List[ZerverFieldsT]:
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:
@@ -630,10 +695,8 @@ def process_uploads(upload_list: List[ZerverFieldsT], upload_dir: str,
logging.info('######### GETTING ATTACHMENTS FINISHED #########\n')
return upload_list
def build_realm_emoji(realm_id: int,
name: str,
id: int,
file_name: str) -> ZerverFieldsT:
def build_realm_emoji(realm_id: int, name: str, id: int, file_name: str) -> ZerverFieldsT:
return model_to_dict(
RealmEmoji(
realm_id=realm_id,
@@ -643,6 +706,7 @@ def build_realm_emoji(realm_id: int,
),
)
def get_emojis(emoji_dir: str, upload: List[str]) -> None:
emoji_url = upload[0]
emoji_path = upload[1]
@@ -653,8 +717,13 @@ def get_emojis(emoji_dir: str, upload: List[str]) -> None:
with open(upload_emoji_path, 'wb') as emoji_file:
shutil.copyfileobj(response.raw, emoji_file)
def process_emojis(zerver_realmemoji: List[ZerverFieldsT], emoji_dir: str,
emoji_url_map: ZerverFieldsT, threads: int) -> List[ZerverFieldsT]:
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:
@@ -670,8 +739,8 @@ def process_emojis(zerver_realmemoji: List[ZerverFieldsT], emoji_dir: str,
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'])
realm_id=emoji['realm'], emoji_file_name=emoji['name']
)
upload_emoji_list.append([emoji_url, emoji_path])
@@ -689,6 +758,7 @@ def process_emojis(zerver_realmemoji: List[ZerverFieldsT], emoji_dir: str,
logging.info('######### GETTING EMOJIS FINISHED #########\n')
return emoji_records
def create_converted_data_files(data: Any, output_dir: str, file_path: str) -> None:
output_file = output_dir + file_path
os.makedirs(os.path.dirname(output_file), exist_ok=True)