diff --git a/confirmation/models.py b/confirmation/models.py index b6a47ae3e3..ce3ea79123 100644 --- a/confirmation/models.py +++ b/confirmation/models.py @@ -16,10 +16,19 @@ from django.http import HttpRequest, HttpResponse from django.shortcuts import render from django.urls import reverse from django.utils.timezone import now as timezone_now +from typing_extensions import Protocol from zerver.models import EmailChangeStatus, MultiuseInvite, PreregistrationUser, Realm, UserProfile +class HasRealmObject(Protocol): + realm: Realm + + +class OptionalHasRealmObject(Protocol): + realm: Optional[Realm] + + class ConfirmationKeyException(Exception): WRONG_LENGTH = 1 EXPIRED = 2 @@ -74,14 +83,16 @@ def get_object_from_key( def create_confirmation_link( - obj: ContentType, confirmation_type: int, url_args: Mapping[str, str] = {} + obj: Union[Realm, HasRealmObject, OptionalHasRealmObject], + confirmation_type: int, + url_args: Mapping[str, str] = {}, ) -> str: key = generate_key() realm = None - if hasattr(obj, "realm"): - realm = obj.realm - elif isinstance(obj, Realm): + if isinstance(obj, Realm): realm = obj + elif hasattr(obj, "realm"): + realm = obj.realm Confirmation.objects.create( content_object=obj, diff --git a/zerver/views/auth.py b/zerver/views/auth.py index 1b6597f2ab..01674ad72d 100644 --- a/zerver/views/auth.py +++ b/zerver/views/auth.py @@ -112,7 +112,7 @@ def create_preregistration_user( password_required: bool = True, full_name: Optional[str] = None, full_name_validated: bool = False, -) -> HttpResponse: +) -> PreregistrationUser: realm = None if not realm_creation: try: