preserve old UX for workspace selector when user has multiple workspaces

daras_ai_v2/base.py

@@ -717,20 +717,17 @@ def _render_publish_form(
     ):
         form_container = gui.div()
 
-        with gui.div(className="my-4"):
-            workspace_create_ref = gui.use_alert_dialog(key="create-workspace-modal")
-            if len(self.request.user.cached_workspaces) == 1:
-                selected_workspace = self.current_workspace
-                col1, col2 = gui.columns([1, 4])
-                with col1:
+        workspace_create_ref = gui.use_alert_dialog(key="create-workspace-modal")
+        if len(self._get_workspace_options()) <= 1:
+            with gui.div(className="my-4 d-flex gap-3"):
+                with gui.div(className="text-nowrap"):
                     gui.write("Workspace")
-                with col2:
+                with gui.div(className="flex-grow-1"):
                     self.render_workspace_author(
-                        workspace=selected_workspace,
+                        workspace=self.current_workspace,
                         image_size="40px",
                         show_as_link=False,
                     )
-
                     with gui.div(
                         className="alert alert-warning mb-0 mt-4 d-flex align-items-baseline gap-2 container-margin-reset"
                     ):
@@ -740,20 +737,24 @@ def _render_publish_form(
                         ):
                             workspace_create_ref.set_open(True)
                         gui.write("to edit with others")
-            else:
+            selected_workspace = self.current_workspace
+            form_buttons_div = gui.div()
+        else:
+            with gui.div(className="mt-4 d-block d-lg-flex justify-content-between"):
                 selected_workspace = self._render_workspace_selector(
                     key="published_run_workspace"
                 )
+                form_buttons_div = gui.div()
 
-            if workspace_create_ref.is_open:
-                render_workspace_create_dialog(
-                    user=self.request.user,
-                    session=self.request.session,
-                    ref=workspace_create_ref,
-                )
+        if workspace_create_ref.is_open:
+            render_workspace_create_dialog(
+                user=self.request.user,
+                session=self.request.session,
+                ref=workspace_create_ref,
+            )
 
         user_can_edit = selected_workspace.id == self.current_pr.workspace_id
-        with gui.div(className="mt-4 mt-lg-0 text-end"):
+        with form_buttons_div, gui.div(className="mt-4 mt-lg-0 text-end"):
             if user_can_edit:
                 pressed_save_as_new = gui.button(
                     f"{icons.fork} Save as New",
@@ -875,39 +876,33 @@ def _render_publish_form(
             )
         raise gui.RedirectException(pr.get_app_url())
 
-    def _render_workspace_selector(self, *, key: str) -> "Workspace":
+    def _get_workspace_options(self) -> dict[int, Workspace]:
         workspace_options = {w.id: w for w in self.request.user.cached_workspaces}
-
         if self.current_pr.workspace_id and self.can_user_edit_published_run(
             self.current_pr
         ):
             # make current_pr.workspace the first option
             workspace_options = {
                 self.current_pr.workspace_id: self.current_pr.workspace
             } | workspace_options
+        return workspace_options
 
+    def _render_workspace_selector(self, *, key: str) -> "Workspace":
+        workspace_options = self._get_workspace_options()
         with gui.div(className="d-flex gap-3"):
             with gui.div(className="mt-2 text-nowrap"):
                 gui.write("Workspace")
 
-            if len(workspace_options) > 1:
-                with gui.div(style=dict(maxWidth="300px", width="100%")):
-                    workspace_id = gui.selectbox(
-                        "",
-                        key=key,
-                        options=workspace_options,
-                        format_func=lambda w_id: workspace_options[w_id].display_html(
-                            self.request.user
-                        ),
-                    )
-                    return workspace_options[workspace_id]
-            else:
-                with gui.div(className="p-2 mb-2"):
-                    self.render_workspace_author(
-                        self.current_workspace,
-                        show_as_link=False,
-                    )
-                return self.current_workspace
+            with gui.div(style=dict(maxWidth="300px", width="100%")):
+                workspace_id = gui.selectbox(
+                    "",
+                    key=key,
+                    options=workspace_options,
+                    format_func=lambda w_id: workspace_options[w_id].display_html(
+                        self.request.user
+                    ),
+                )
+                return workspace_options[workspace_id]
 
     def _get_default_pr_title(self):
         return f"{self.request.user.first_name_possesive()} {self.get_run_title(self.current_sr, self.current_pr)}"