Fix ImportError after refactoring (PY-21487)