From 1ca3ef14a300d9d19b4fbd5fb88d20a3e1521ac4 Mon Sep 17 00:00:00 2001
From: skshetry <18718008+skshetry@users.noreply.github.com>
Date: Thu, 30 Mar 2023 08:27:47 +0545
Subject: [PATCH] exp-push: show url to the Studio project (#9278)

---
 dvc/commands/experiments/push.py      |  9 +++++---
 dvc/repo/experiments/push.py          | 31 +++++++++++++++++++--------
 tests/func/experiments/test_remote.py |  5 ++++-
 3 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/dvc/commands/experiments/push.py b/dvc/commands/experiments/push.py
index 66631bcbd..239971a37 100644
--- a/dvc/commands/experiments/push.py
+++ b/dvc/commands/experiments/push.py
@@ -23,7 +23,7 @@ class CmdExperimentsPush(CmdBase):
 
         self.raise_error_if_all_disabled()
 
-        pushed_exps = self.repo.experiments.push(
+        result = self.repo.experiments.push(
             self.args.git_remote,
             self.args.experiment,
             all_commits=self.args.all_commits,
@@ -35,14 +35,17 @@ class CmdExperimentsPush(CmdBase):
             jobs=self.args.jobs,
             run_cache=self.args.run_cache,
         )
-
-        if pushed_exps:
+        if pushed_exps := result.get("pushed"):
             ui.write(
                 f"Pushed experiment {humanize.join(map(repr, pushed_exps))} "
                 f"to Git remote {self.args.git_remote!r}."
             )
         else:
             ui.write("No experiments to push.")
+
+        if project_url := result.get("url"):
+            ui.write("[yellow]View your experiments at", project_url, styled=True)
+
         if not self.args.push_cache:
             ui.write(
                 "To push cached outputs",
diff --git a/dvc/repo/experiments/push.py b/dvc/repo/experiments/push.py
index 27e25987d..25121293c 100644
--- a/dvc/repo/experiments/push.py
+++ b/dvc/repo/experiments/push.py
@@ -1,5 +1,15 @@
 import logging
-from typing import TYPE_CHECKING, Any, Iterable, List, Mapping, Optional, Set, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Set,
+    Union,
+)
 
 from funcy import compact, group_by
 from scmrepo.git.backend.base import SyncStatus
@@ -20,25 +30,27 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-def notify_refs_to_studio(repo: "Repo", git_remote: str, **refs: List[str]) -> None:
+def notify_refs_to_studio(
+    repo: "Repo", git_remote: str, **refs: List[str]
+) -> Optional[str]:
     config = repo.config["feature"]
     refs = compact(refs)
     if not refs or env2bool("DVC_TEST"):
-        return
+        return None
 
     if not (config.get("studio_token") or config["push_exp_to_studio"]):
         logger.debug(
             "Either feature.studio_token or feature.push_exp_to_studio config "
             "needs to be set."
         )
-        return
+        return None
 
     import os
 
     token = os.environ.get("STUDIO_TOKEN") or config.get("studio_token")
     if not token:
         logger.debug("Studio token not found.")
-        return
+        return None
 
     from dulwich.porcelain import get_remote_repo
 
@@ -46,7 +58,8 @@ def notify_refs_to_studio(repo: "Repo", git_remote: str, **refs: List[str]) -> N
 
     _, repo_url = get_remote_repo(repo.scm.dulwich.repo, git_remote)
     studio_url = config.get("studio_url")
-    studio.notify_refs(repo_url, token, studio_url=studio_url, **refs)
+    d = studio.notify_refs(repo_url, token, studio_url=studio_url, **refs)
+    return d.get("url")
 
 
 @locked
@@ -61,7 +74,7 @@ def push(  # noqa: C901
     force: bool = False,
     push_cache: bool = False,
     **kwargs: Any,
-) -> Iterable[str]:
+) -> Dict[str, Any]:
     exp_ref_set: Set["ExpRefInfo"] = set()
     assert isinstance(repo.scm, Git)
     if all_commits:
@@ -106,8 +119,8 @@ def push(  # noqa: C901
 
     refs = push_result[SyncStatus.SUCCESS]
     pushed_refs = [str(r) for r in refs]
-    notify_refs_to_studio(repo, git_remote, pushed=pushed_refs)
-    return [ref.name for ref in refs]
+    url = notify_refs_to_studio(repo, git_remote, pushed=pushed_refs)
+    return {"pushed": [ref.name for ref in refs], "url": url}
 
 
 def _push(
diff --git a/tests/func/experiments/test_remote.py b/tests/func/experiments/test_remote.py
index 5251cb71c..c9df12ca0 100644
--- a/tests/func/experiments/test_remote.py
+++ b/tests/func/experiments/test_remote.py
@@ -76,7 +76,10 @@ def test_push_diverged(tmp_dir, scm, dvc, git_upstream, exp_stage):
 
     git_upstream.tmp_dir.scm.set_ref(str(ref_info), remote_rev)
 
-    assert dvc.experiments.push(git_upstream.remote, [ref_info.name]) == []
+    assert dvc.experiments.push(git_upstream.remote, [ref_info.name]) == {
+        "pushed": [],
+        "url": None,
+    }
     assert git_upstream.tmp_dir.scm.get_ref(str(ref_info)) == remote_rev
 
     dvc.experiments.push(git_upstream.remote, [ref_info.name], force=True)