def mock_subclasses(baseclass, *subclasses):
    """Mocks baseclass so that it only has the given child subclasses."""
    with ExitStack() as stack:
        mgr = mock.patch.object(baseclass, "__subclasses__", return_value=[*subclasses])
        stack.enter_context(mgr)
        for mgr in [mock.patch.object(s, "__subclasses__", return_value=[]) for s in subclasses]:
            stack.enter_context(mgr)
        yield None
Mocks baseclass so that it only has the given child subclasses.
tests/tests_pytorch/test_cli.py