From 35418de9dfd7b0d4f68e5868c94e3f6c6d58191f Mon Sep 17 00:00:00 2001
From: Julien Chaumond <julien@huggingface.co>
Date: Tue, 2 Mar 2021 15:26:02 +0100
Subject: [PATCH] Helper: when pushing, return commit url (#20)

cc @philschmid
---
 src/huggingface_hub/repository.py | 56 +++++++++++++++++++++++++++++--
 tests/test_repository.py          |  7 +++-
 2 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/src/huggingface_hub/repository.py b/src/huggingface_hub/repository.py
index 0aad5b9a1e..b435ff2a90 100644
--- a/src/huggingface_hub/repository.py
+++ b/src/huggingface_hub/repository.py
@@ -222,6 +222,52 @@ def git_config_username_and_email(
         except subprocess.CalledProcessError as exc:
             raise EnvironmentError(exc.stderr)
 
+    def git_head_hash(self) -> str:
+        """
+        Get commit sha on top of HEAD.
+        """
+        try:
+            p = subprocess.run(
+                "git rev-parse HEAD".split(),
+                stderr=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                encoding="utf-8",
+                check=True,
+                cwd=self.local_dir,
+            )
+            return p.stdout.strip()
+        except subprocess.CalledProcessError as exc:
+            raise EnvironmentError(exc.stderr)
+
+    def git_remote_url(self) -> str:
+        """
+        Get URL to origin remote.
+        """
+        try:
+            p = subprocess.run(
+                "git config --get remote.origin.url".split(),
+                stderr=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                encoding="utf-8",
+                check=True,
+                cwd=self.local_dir,
+            )
+            return p.stdout.strip()
+        except subprocess.CalledProcessError as exc:
+            raise EnvironmentError(exc.stderr)
+
+    def git_head_commit_url(self) -> str:
+        """
+        Get URL to last commit on HEAD
+        We assume it's been pushed, and the url scheme is
+        the same one as for GitHub or HuggingFace.
+        """
+        sha = self.git_head_hash()
+        url = self.git_remote_url()
+        if url.endswith("/"):
+            url = url[:-1]
+        return f"{url}/commit/{sha}"
+
     def lfs_track(self, patterns: Union[str, List[str]]):
         """
         Tell git-lfs to track those files.
@@ -319,9 +365,11 @@ def git_commit(self, commit_message="commit files to HF hub"):
             else:
                 raise EnvironmentError(exc.stdout)
 
-    def git_push(self):
+    def git_push(self) -> str:
         """
         git push
+
+        Returns url to commit on remote repo.
         """
         try:
             result = subprocess.run(
@@ -336,7 +384,9 @@ def git_push(self):
         except subprocess.CalledProcessError as exc:
             raise EnvironmentError(exc.stderr)
 
-    def push_to_hub(self, commit_message="commit files to HF hub"):
+        return self.git_head_commit_url()
+
+    def push_to_hub(self, commit_message="commit files to HF hub") -> str:
         """
         Helper to add, commit, and pushe file to remote repository on the HuggingFace Hub.
         Args:
@@ -344,4 +394,4 @@ def push_to_hub(self, commit_message="commit files to HF hub"):
         """
         self.git_add()
         self.git_commit(commit_message)
-        self.git_push()
+        return self.git_push()
diff --git a/tests/test_repository.py b/tests/test_repository.py
index 37e70f24b5..8de6cf3543 100644
--- a/tests/test_repository.py
+++ b/tests/test_repository.py
@@ -19,6 +19,7 @@
 import time
 import unittest
 
+import requests
 from huggingface_hub.hf_api import HfApi
 from huggingface_hub.repository import Repository
 
@@ -110,7 +111,11 @@ def test_add_commit_push(self):
         repo.git_add()
         repo.git_commit()
         try:
-            repo.git_push()
+            url = repo.git_push()
         except subprocess.CalledProcessError as exc:
             print(exc.stderr)
             raise exc
+        # Check that the returned commit url
+        # actually exists.
+        r = requests.head(url)
+        r.raise_for_status()