Git prune (#450)

* Git lfs prune

* Tests

* Add it to push and context manager

* Adjust documentation

Author: LysandreJik

src/huggingface_hub/repository.py

@@ -29,20 +29,27 @@ def __init__(
         is_done_method: Callable,
         status_method: Callable,
         process: subprocess.Popen,
+        post_method: Optional[Callable] = None,
     ):
         self.title = title
         self._is_done = is_done_method
         self._status = status_method
         self._process = process
         self._stderr = ""
         self._stdout = ""
+        self._post_method = post_method
 
     @property
     def is_done(self) -> bool:
         """
         Whether the process is done.
         """
-        return self._is_done()
+        result = self._is_done()
+
+        if result and self._post_method is not None:
+            self._post_method()
+
+        return result
 
     @property
     def status(self) -> int:
@@ -840,6 +847,27 @@ def auto_track_large_files(self, pattern: Optional[str] = ".") -> List[str]:
 
         return files_to_be_tracked_with_lfs
 
+    def lfs_prune(self, recent=False):
+        """
+        git lfs prune
+        """
+        args = "git lfs prune".split()
+        if recent:
+            args.append("--recent")
+        try:
+            with lfs_log_progress():
+                result = subprocess.run(
+                    args,
+                    stderr=subprocess.PIPE,
+                    stdout=subprocess.PIPE,
+                    check=True,
+                    encoding="utf-8",
+                    cwd=self.local_dir,
+                )
+                logger.info(result.stdout)
+        except subprocess.CalledProcessError as exc:
+            raise EnvironmentError(exc.stderr)
+
     def git_pull(self, rebase: Optional[bool] = False):
         """
         git pull
@@ -914,6 +942,7 @@ def git_push(
         self,
         upstream: Optional[str] = None,
         blocking: Optional[bool] = True,
+        auto_lfs_prune: Optional[bool] = False,
     ) -> Union[str, Tuple[str, CommandInProgress]]:
         """
         git push
@@ -931,6 +960,8 @@ def git_push(
                 Setting this to `False` will return an `CommandInProgress` object
                 which has an `is_done` property. This property will be set to
                 `True` when the push is finished.
+            auto_lfs_prune (`bool`, defaults to `False`):
+                Whether to automatically prune files once they have been pushed to the remote.
         """
         command = "git push"
 
@@ -986,12 +1017,16 @@ def status_method():
                 is_done_method=lambda: process.poll() is not None,
                 status_method=status_method,
                 process=process,
+                post_method=self.lfs_prune if auto_lfs_prune else None,
             )
 
             self.command_queue.append(command)
 
             return self.git_head_commit_url(), command
 
+        if auto_lfs_prune:
+            self.lfs_prune()
+
         return self.git_head_commit_url()
 
     def git_checkout(self, revision: str, create_branch_ok: Optional[bool] = False):
@@ -1172,6 +1207,7 @@ def push_to_hub(
         commit_message: Optional[str] = "commit files to HF hub",
         blocking: Optional[bool] = True,
         clean_ok: Optional[bool] = False,
+        auto_lfs_prune: Optional[bool] = False,
     ) -> Optional[str]:
         """
         Helper to add, commit, and push files to remote repository on the HuggingFace Hub.
@@ -1185,14 +1221,18 @@ def push_to_hub(
             clean_ok (`bool`, `optional`, defaults to `False`):
                 If True, this function will return None if the repo is untouched.
                 Default behavior is to fail because the git command fails.
+            auto_lfs_prune (`bool`, defaults to `False`):
+                Whether to automatically prune files once they have been pushed to the remote.
         """
         if clean_ok and self.is_repo_clean():
             logger.info("Repo currently clean.  Ignoring push_to_hub")
             return None
         self.git_add(auto_lfs_track=True)
         self.git_commit(commit_message)
         return self.git_push(
-            upstream=f"origin {self.current_branch}", blocking=blocking
+            upstream=f"origin {self.current_branch}",
+            blocking=blocking,
+            auto_lfs_prune=auto_lfs_prune,
         )
 
     @contextmanager
@@ -1202,6 +1242,7 @@ def commit(
         branch: Optional[str] = None,
         track_large_files: Optional[bool] = True,
         blocking: Optional[bool] = True,
+        auto_lfs_prune: Optional[bool] = False,
     ):
         """
         Context manager utility to handle committing to a repository. This automatically tracks large files (>10Mb)
@@ -1216,6 +1257,8 @@ def commit(
                 Whether to automatically track large files or not. Will do so by default.
             blocking (`bool`, `optional`, defaults to `True`):
                 Whether the function should return only when the `git push` has finished.
+            auto_lfs_prune (`bool`, defaults to `True`):
+                Whether to automatically prune files once they have been pushed to the remote.
 
         Examples:
 
@@ -1270,7 +1313,9 @@ def commit(
 
             try:
                 self.git_push(
-                    upstream=f"origin {self.current_branch}", blocking=blocking
+                    upstream=f"origin {self.current_branch}",
+                    blocking=blocking,
+                    auto_lfs_prune=auto_lfs_prune,
                 )
             except OSError as e:
                 # If no changes are detected, there is nothing to commit.

tests/test_repository.py

@@ -735,6 +735,168 @@ def test_delete_tag(self):
         repo.delete_tag("v4.6.0", remote="origin")
         self.assertFalse(repo.tag_exists("v4.6.0", remote="origin"))
 
+    def test_lfs_prune(self):
+        repo = Repository(
+            WORKING_REPO_DIR,
+            clone_from=f"{USER}/{REPO_NAME}",
+            use_auth_token=self._token,
+            git_user="ci",
+            git_email="[email protected]",
+            revision="main",
+        )
+
+        with repo.commit("Committing LFS file"):
+            with open("file.bin", "w+") as f:
+                f.write("Random string 1")
+
+        with repo.commit("Committing LFS file"):
+            with open("file.bin", "w+") as f:
+                f.write("Random string 2")
+
+        root_directory = pathlib.Path(repo.local_dir) / ".git" / "lfs"
+        git_lfs_files_size = sum(
+            f.stat().st_size for f in root_directory.glob("**/*") if f.is_file()
+        )
+        repo.lfs_prune()
+        post_prune_git_lfs_files_size = sum(
+            f.stat().st_size for f in root_directory.glob("**/*") if f.is_file()
+        )
+
+        # Size of the directory holding LFS files was reduced
+        self.assertLess(post_prune_git_lfs_files_size, git_lfs_files_size)
+
+    def test_lfs_prune_git_push(self):
+        repo = Repository(
+            WORKING_REPO_DIR,
+            clone_from=f"{USER}/{REPO_NAME}",
+            use_auth_token=self._token,
+            git_user="ci",
+            git_email="[email protected]",
+            revision="main",
+        )
+
+        with repo.commit("Committing LFS file"):
+            with open("file.bin", "w+") as f:
+                f.write("Random string 1")
+
+        root_directory = pathlib.Path(repo.local_dir) / ".git" / "lfs"
+        git_lfs_files_size = sum(
+            f.stat().st_size for f in root_directory.glob("**/*") if f.is_file()
+        )
+
+        with open(os.path.join(repo.local_dir, "file.bin"), "w+") as f:
+            f.write("Random string 2")
+
+        repo.git_add()
+        repo.git_commit("New commit")
+        repo.git_push(auto_lfs_prune=True)
+
+        post_prune_git_lfs_files_size = sum(
+            f.stat().st_size for f in root_directory.glob("**/*") if f.is_file()
+        )
+
+        # Size of the directory holding LFS files is the exact same
+        self.assertEqual(post_prune_git_lfs_files_size, git_lfs_files_size)
+
+    def test_lfs_prune_git_push_non_blocking(self):
+        repo = Repository(
+            WORKING_REPO_DIR,
+            clone_from=f"{USER}/{REPO_NAME}",
+            use_auth_token=self._token,
+            git_user="ci",
+            git_email="[email protected]",
+            revision="main",
+        )
+
+        with repo.commit("Committing LFS file"):
+            with open("file.bin", "w+") as f:
+                f.write("Random string 1")
+
+        root_directory = pathlib.Path(repo.local_dir) / ".git" / "lfs"
+        git_lfs_files_size = sum(
+            f.stat().st_size for f in root_directory.glob("**/*") if f.is_file()
+        )
+
+        with open(os.path.join(repo.local_dir, "file.bin"), "w+") as f:
+            f.write("Random string 2")
+
+        repo.git_add()
+        repo.git_commit("New commit")
+        repo.git_push(blocking=False, auto_lfs_prune=True)
+
+        while len(repo.commands_in_progress):
+            time.sleep(0.2)
+
+        post_prune_git_lfs_files_size = sum(
+            f.stat().st_size for f in root_directory.glob("**/*") if f.is_file()
+        )
+
+        # Size of the directory holding LFS files is the exact same
+        self.assertEqual(post_prune_git_lfs_files_size, git_lfs_files_size)
+
+    def test_lfs_prune_context_manager(self):
+        repo = Repository(
+            WORKING_REPO_DIR,
+            clone_from=f"{USER}/{REPO_NAME}",
+            use_auth_token=self._token,
+            git_user="ci",
+            git_email="[email protected]",
+            revision="main",
+        )
+
+        with repo.commit("Committing LFS file"):
+            with open("file.bin", "w+") as f:
+                f.write("Random string 1")
+
+        root_directory = pathlib.Path(repo.local_dir) / ".git" / "lfs"
+        git_lfs_files_size = sum(
+            f.stat().st_size for f in root_directory.glob("**/*") if f.is_file()
+        )
+
+        with repo.commit("Committing LFS file", auto_lfs_prune=True):
+            with open("file.bin", "w+") as f:
+                f.write("Random string 2")
+
+        post_prune_git_lfs_files_size = sum(
+            f.stat().st_size for f in root_directory.glob("**/*") if f.is_file()
+        )
+
+        # Size of the directory holding LFS files is the exact same
+        self.assertEqual(post_prune_git_lfs_files_size, git_lfs_files_size)
+
+    def test_lfs_prune_context_manager_non_blocking(self):
+        repo = Repository(
+            WORKING_REPO_DIR,
+            clone_from=f"{USER}/{REPO_NAME}",
+            use_auth_token=self._token,
+            git_user="ci",
+            git_email="[email protected]",
+            revision="main",
+        )
+
+        with repo.commit("Committing LFS file"):
+            with open("file.bin", "w+") as f:
+                f.write("Random string 1")
+
+        root_directory = pathlib.Path(repo.local_dir) / ".git" / "lfs"
+        git_lfs_files_size = sum(
+            f.stat().st_size for f in root_directory.glob("**/*") if f.is_file()
+        )
+
+        with repo.commit("Committing LFS file", auto_lfs_prune=True, blocking=False):
+            with open("file.bin", "w+") as f:
+                f.write("Random string 2")
+
+        while len(repo.commands_in_progress):
+            time.sleep(0.2)
+
+        post_prune_git_lfs_files_size = sum(
+            f.stat().st_size for f in root_directory.glob("**/*") if f.is_file()
+        )
+
+        # Size of the directory holding LFS files is the exact same
+        self.assertEqual(post_prune_git_lfs_files_size, git_lfs_files_size)
+
 
 class RepositoryOfflineTest(RepositoryCommonTest):
     @classmethod