Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add url attribute to Collection class #1695

Merged
merged 2 commits into from
Sep 25, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion docs/source/en/guides/collections.md
Original file line number Diff line number Diff line change
Expand Up @@ -79,11 +79,13 @@ It will return a [`Collection`] object with the high-level metadata (title, desc

```py
>>> collection.slug
'iccv-2023-15e23b46cb98efca45'
'owner/iccv-2023-15e23b46cb98efca45'
>>> collection.title
"ICCV 2023"
>>> collection.owner
"username"
>>> collection.url
'https://huggingface.co/collections/owner/iccv-2023-15e23b46cb98efca45'
```

## Manage items in a collection
Expand Down
17 changes: 12 additions & 5 deletions src/huggingface_hub/hf_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -663,6 +663,8 @@ class Collection(ReprMixin):
Whether the collection is private or not.
theme (`str`):
Theme of the collection. E.g. `"green"`.
url (`str`):
URL for the collection on the Hub.
"""

slug: str
Expand All @@ -676,7 +678,7 @@ class Collection(ReprMixin):
private: bool
theme: str

def __init__(self, data: Dict) -> None:
def __init__(self, data: Dict, endpoint: Optional[str] = None) -> None:
# Collection info
self.slug = data["slug"]
self.title = data["title"]
Expand All @@ -690,6 +692,11 @@ def __init__(self, data: Dict) -> None:
self.position = data["position"]
self.theme = data["theme"]

# (internal)
if endpoint is None:
endpoint = ENDPOINT
self.url = f"{ENDPOINT}/collections/{self.slug}"


class ModelSearchArguments(AttributeDictionary):
"""
Expand Down Expand Up @@ -5840,7 +5847,7 @@ def get_collection(self, collection_slug: str, *, token: Optional[str] = None) -
f"{self.endpoint}/api/collections/{collection_slug}", headers=self._build_hf_headers(token=token)
)
hf_raise_for_status(r)
return Collection(r.json())
return Collection(r.json(), endpoint=self.endpoint)

def create_collection(
self,
Expand Down Expand Up @@ -5905,7 +5912,7 @@ def create_collection(
return self.get_collection(slug, token=token)
else:
raise
return Collection(r.json())
return Collection(r.json(), endpoint=self.endpoint)

def update_collection_metadata(
self,
Expand Down Expand Up @@ -5970,7 +5977,7 @@ def update_collection_metadata(
json={key: value for key, value in payload.items() if value is not None},
)
hf_raise_for_status(r)
return Collection(r.json()["data"])
return Collection(r.json()["data"], endpoint=self.endpoint)

def delete_collection(
self, collection_slug: str, *, missing_ok: bool = False, token: Optional[str] = None
Expand Down Expand Up @@ -6078,7 +6085,7 @@ def add_collection_item(
return self.get_collection(collection_slug, token=token)
else:
raise
return Collection(r.json())
return Collection(r.json(), endpoint=self.endpoint)

def update_collection_item(
self,
Expand Down
1 change: 1 addition & 0 deletions tests/test_hf_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -3225,6 +3225,7 @@ def test_create_collection_with_description(self) -> None:
self.assertEqual(collection.description, "Contains a lot of cool stuff")
self.assertEqual(collection.items, [])
self.assertTrue(collection.slug.startswith(self.slug_prefix))
self.assertEqual(collection.url, f"{ENDPOINT_STAGING}/collections/{collection.slug}")

def test_create_collection_exists_ok(self) -> None:
# Create collection once without description
Expand Down