diff --git a/bots/models.py b/bots/models.py index 31ee78cf5..3e1a61dbf 100644 --- a/bots/models.py +++ b/bots/models.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import datetime import typing from multiprocessing.pool import ThreadPool @@ -40,14 +42,14 @@ class PublishedRunVisibility(models.IntegerChoices): INTERNAL = 3 @classmethod - def choices_for_workspace( - cls, workspace: typing.Optional["Workspace"] - ) -> typing.Iterable["PublishedRunVisibility"]: - if not workspace or workspace.is_personal: - return [cls.UNLISTED, cls.PUBLIC] + def choices_for_pr( + cls, pr: PublishedRun + ) -> typing.Iterable[PublishedRunVisibility]: + if not pr.workspace or pr.workspace.is_personal: + return {cls.UNLISTED, cls.PUBLIC, PublishedRunVisibility(pr.visibility)} else: # TODO: Add cls.PUBLIC when team-handles are added - return [cls.UNLISTED, cls.INTERNAL] + return {cls.UNLISTED, cls.INTERNAL, PublishedRunVisibility(pr.visibility)} @classmethod def get_default_for_workspace(cls, workspace: typing.Optional["Workspace"]): diff --git a/daras_ai_v2/base.py b/daras_ai_v2/base.py index e07a1bd9f..3a806bb13 100644 --- a/daras_ai_v2/base.py +++ b/daras_ai_v2/base.py @@ -557,9 +557,7 @@ def _render_share_modal(self, dialog: gui.AlertDialogRef): options = { str(enum.value): enum.help_text(self.current_pr.workspace) - for enum in PublishedRunVisibility.choices_for_workspace( - self.current_pr.workspace - ) + for enum in PublishedRunVisibility.choices_for_pr(self.current_pr) } published_run_visibility = PublishedRunVisibility( int(