Notes

I could just use the inbuilt notebook with .py files, as suggested here:

Exploring an Alternative to Jupyter Notebooks for Python Development

There still isn't any TOC functionality, but there will be collapsable headers and compatibility with other VS Code extensions as described here:

Medium

VS Code can also be used from inside Google Colab:

VSCode on Google Colab

as well as other cloud providers:

Unleash the power of Visual Studio Code (VSCode) on Google Cloud Platform Virtual Machine

Conclusion

For now, I'll focus on using VS Code, for its intelligent autocomplete, debugging and extensions, which will even improve on the next release.