Have you tried pyCharm? I believe even community (free) edition should have support You are looking for
If it's for student, you could maybe get educational license, that is completely free for all their products, but I imagine it could be quite a hassle
The VS Code implementation of notebooks has too much vertical space for me.