From 0c6924de39e949f1e6e24013af18857d22829ed0 Mon Sep 17 00:00:00 2001 From: Quentin Lhoest Date: Wed, 29 Sep 2021 17:06:49 +0200 Subject: [PATCH] add timeout to dataset_info --- src/huggingface_hub/hf_api.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/huggingface_hub/hf_api.py b/src/huggingface_hub/hf_api.py index 6556f6f146..6f77630162 100644 --- a/src/huggingface_hub/hf_api.py +++ b/src/huggingface_hub/hf_api.py @@ -492,7 +492,11 @@ def list_datasets( return [DatasetInfo(**x) for x in d] def model_info( - self, repo_id: str, revision: Optional[str] = None, token: Optional[str] = None + self, + repo_id: str, + revision: Optional[str] = None, + token: Optional[str] = None, + timeout: Optional[float] = None, ) -> ModelInfo: """ Get info on one specific model on huggingface.co @@ -509,7 +513,7 @@ def model_info( headers = ( {"authorization": "Bearer {}".format(token)} if token is not None else None ) - r = requests.get(path, headers=headers) + r = requests.get(path, headers=headers, timeout=timeout) r.raise_for_status() d = r.json() return ModelInfo(**d) @@ -532,7 +536,11 @@ def list_repos_objs( return [RepoObj(**x) for x in d] def dataset_info( - self, repo_id: str, revision: Optional[str] = None, token: Optional[str] = None + self, + repo_id: str, + revision: Optional[str] = None, + token: Optional[str] = None, + timeout: Optional[float] = None, ) -> DatasetInfo: """ Get info on one specific dataset on huggingface.co @@ -550,7 +558,7 @@ def dataset_info( {"authorization": "Bearer {}".format(token)} if token is not None else None ) params = {"full": "true"} - r = requests.get(path, headers=headers, params=params) + r = requests.get(path, headers=headers, params=params, timeout=timeout) r.raise_for_status() d = r.json() return DatasetInfo(**d)