From 396d6da37b424fbc5e27ee56356115be48f3704d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?=
 =?UTF-8?q?=E0=A4=A4=29?= <suagatchhetri@outlook.com>
Date: Sat, 30 Dec 2023 14:29:04 +0545
Subject: [PATCH] storages: add rich repr

---
 src/dvc_data/index/index.py | 38 ++++++++++++++++++++++---------------
 1 file changed, 23 insertions(+), 15 deletions(-)

diff --git a/src/dvc_data/index/index.py b/src/dvc_data/index/index.py
index acf8ccf8..068e3dfc 100644
--- a/src/dvc_data/index/index.py
+++ b/src/dvc_data/index/index.py
@@ -184,11 +184,14 @@ def __init__(
         super().__init__(key, read_only=read_only)
 
     def __repr__(self) -> str:
-        key = self.key
-        odb = self.odb
-        index = self.index
-        read_only = self.read_only
-        return f"{type(self).__name__}({key=!r}, {odb=!r}, {index=!r}, {read_only=!r})"
+        params = ", ".join(f"{k}={v!r}" for k, v, *_ in self.__rich_repr__())
+        return f"{self.__class__.__name__}({params})"
+
+    def __rich_repr__(self):
+        yield "key", self.key
+        yield "odb", self.odb
+        yield "index", self.index, None
+        yield "read_only", self.read_only, False
 
     @property
     def fs(self):
@@ -253,16 +256,15 @@ def __init__(
         super().__init__(key, read_only=read_only)
 
     def __repr__(self) -> str:
-        key = self.key
-        fs = self.fs
-        path = self.path
-        index = self.index
-        prefix = self.prefix
-        read_only = self.read_only
-        return (
-            f"{self.__class__.__name__}("
-            f"{key=!r}, {fs=!r}, {path=!r}, {index=!r}, {prefix=!r}, {read_only=!r})"
-        )
+        params = ", ".join(f"{k}={v!r}" for k, v, *_ in self.__rich_repr__())
+        return f"{self.__class__.__name__}({params})"
+
+    def __rich_repr__(self):
+        yield "key", self.key
+        yield "fs", self.fs
+        yield "index", self.index, None
+        yield "prefix", self.prefix, None
+        yield "read_only", self.read_only, False
 
     @property
     def fs(self):
@@ -340,6 +342,12 @@ class StorageMapping(MutableMapping):
     def __init__(self, *args, **kwargs):
         self._map = dict(*args, **kwargs)
 
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self._map!r})"
+
+    def __rich_repr__(self):
+        yield self._map
+
     def __setitem__(self, key, value):
         self._map[key] = value