Improve type annotations
Adds correct type annotations for lists and dicts in function definitions. I also included some local type annotations where things are a bit complicated.
Watch https://www.youtube.com/watch?v=pMgmKJyWKn8 (also a nice explainer on how tests and type annotations play well together!) and https://www.youtube.com/watch?v=SMXsIX3PZ5w&t=495s for details.
There is some weak requirement of Python>3.9 for this to play nicely.