From e16dacf7e37c4f1b7c4443f7c36009a2159bccf9 Mon Sep 17 00:00:00 2001
From: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>
Date: Tue, 21 Feb 2023 13:18:11 +0800
Subject: [PATCH] Refine the repr of Registry (#942)

* Refine the repr of Registry

* fix ut

* fix ut
---
 mmengine/registry/registry.py        |  9 ++++-----
 tests/test_registry/test_registry.py | 11 +++--------
 2 files changed, 7 insertions(+), 13 deletions(-)

diff --git a/mmengine/registry/registry.py b/mmengine/registry/registry.py
index a8d385db..19693a13 100644
--- a/mmengine/registry/registry.py
+++ b/mmengine/registry/registry.py
@@ -1,6 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import inspect
-import io
 import logging
 import sys
 from collections.abc import Callable
@@ -131,11 +130,11 @@ class Registry:
         for name, obj in sorted(self._module_dict.items()):
             table.add_row(name, str(obj))
 
-        with io.StringIO() as sio:
-            console = Console(file=sio)
+        console = Console()
+        with console.capture() as capture:
             console.print(table, end='')
-            table_str = sio.getvalue()
-        return table_str
+
+        return capture.get()
 
     @staticmethod
     def infer_scope() -> str:
diff --git a/tests/test_registry/test_registry.py b/tests/test_registry/test_registry.py
index 507cda00..510d9141 100644
--- a/tests/test_registry/test_registry.py
+++ b/tests/test_registry/test_registry.py
@@ -1,9 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-import io
 import time
 
 import pytest
-from rich.console import Console
 
 from mmengine.config import Config, ConfigDict  # type: ignore
 from mmengine.registry import (DefaultScope, Registry, build_from_cfg,
@@ -475,12 +473,9 @@ class TestRegistry:
         class Munchkin:
             pass
 
-        with io.StringIO() as sio:
-            console = Console(file=sio)
-            console.print(CATS, end='')
-            repr_str = sio.getvalue()
-
-        assert repr(CATS) == repr_str
+        assert 'Registry of cat' in repr(CATS)
+        assert 'BritishShorthair' in repr(CATS)
+        assert 'Munchkin' in repr(CATS)
 
 
 @pytest.mark.parametrize('cfg_type', [dict, ConfigDict, Config])