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])