From e4679cddced7d746427066a78e8079fb40e51528 Mon Sep 17 00:00:00 2001
From: Julien Chaumond <chaumond@gmail.com>
Date: Fri, 6 Dec 2019 11:56:23 -0500
Subject: [PATCH] [cli] Uploads: add progress bar (#2078)

* [cli] Uploads: add progress bar

see https://github.com/huggingface/transformers/pull/2044#discussion_r354057827 for context

* rename + documentation

* Add auto-referential comment
---
 transformers/hf_api.py | 34 +++++++++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/transformers/hf_api.py b/transformers/hf_api.py
index c21592a838..3bbb6c567a 100644
--- a/transformers/hf_api.py
+++ b/transformers/hf_api.py
@@ -16,10 +16,11 @@ from __future__ import absolute_import, division, print_function
 
 import os
 from os.path import expanduser
-import six
 
 import requests
+import six
 from requests.exceptions import HTTPError
+from tqdm import tqdm
 
 ENDPOINT = "https://huggingface.co"
 
@@ -129,10 +130,13 @@ class HfApi:
         # Even though we presign with the correct content-type,
         # the client still has to specify it when uploading the file.
         with open(filepath, "rb") as f:
+            pf = TqdmProgressFileReader(f)
+
             r = requests.put(urls.write, data=f, headers={
                 "content-type": urls.type,
             })
             r.raise_for_status()
+            pf.close()
         return urls.access
 
     def list_objs(self, token):
@@ -148,6 +152,34 @@ class HfApi:
 
 
 
+class TqdmProgressFileReader:
+    """
+    Wrap an io.BufferedReader `f` (such as the output of `open(…, "rb")`)
+    and override `f.read()` so as to display a tqdm progress bar.
+
+    see github.com/huggingface/transformers/pull/2078#discussion_r354739608
+    for implementation details.
+    """
+    def __init__(
+        self,
+        f   # type: io.BufferedReader
+    ):
+        self.f = f
+        self.total_size = os.fstat(f.fileno()).st_size # type: int
+        self.pbar = tqdm(total=self.total_size, leave=False)
+        if six.PY3:
+            # does not work unless PY3
+            # no big deal as the CLI does not currently support PY2 anyways.
+            self.read = f.read
+            f.read = self._read
+
+    def _read(self, n=-1):
+        self.pbar.update(n)
+        return self.read(n)
+
+    def close(self):
+        self.pbar.close()
+
 
 
 class HfFolder: