I haven't used it much, but I believe you can do this by using (or more likely slightly extending) the init_cell nbextension.

https://jupyter-contrib-nbextensions.readthedocs.io/en/lates...