From 79e94fcdf3d4378ddcdf7e130bb1ae23d99c6fce Mon Sep 17 00:00:00 2001
From: Quentin Lhoest <42851186+lhoestq@users.noreply.github.com>
Date: Fri, 8 Dec 2023 14:26:53 +0100
Subject: [PATCH] Add IterableDataset `__repr__` (#6480)

add iterabledataset repr
---
 src/datasets/dataset_dict.py     | 5 +++++
 src/datasets/iterable_dataset.py | 3 +++
 2 files changed, 8 insertions(+)

diff --git a/src/datasets/dataset_dict.py b/src/datasets/dataset_dict.py
index 9b06bf4a83a..e955a64a587 100644
--- a/src/datasets/dataset_dict.py
+++ b/src/datasets/dataset_dict.py
@@ -1815,6 +1815,11 @@ def push_to_hub(
 
 
 class IterableDatasetDict(dict):
+    def __repr__(self):
+        repr = "\n".join([f"{k}: {v}" for k, v in self.items()])
+        repr = re.sub(r"^", " " * 4, repr, 0, re.M)
+        return f"IterableDatasetDict({{\n{repr}\n}})"
+
     def with_format(
         self,
         type: Optional[str] = None,
diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py
index d9c6057929c..ccc4c77b2ec 100644
--- a/src/datasets/iterable_dataset.py
+++ b/src/datasets/iterable_dataset.py
@@ -1220,6 +1220,9 @@ def __init__(
         self._token_per_repo_id: Dict[str, Union[str, bool, None]] = token_per_repo_id or {}
         _maybe_add_torch_iterable_dataset_parent_class(self.__class__)
 
+    def __repr__(self):
+        return f"IterableDataset({{\n    features: {list(self._info.features.keys()) if self._info.features is not None else 'Unknown'},\n    n_shards: {self.n_shards}\n}})"
+
     def __getstate__(self):
         return self.__dict__