Clean up code

[ismael-mendoza]

• [x] #71
• [ ] commented out code
• [x] places where black formatting makes code particularly difficult to read (refactor into multiple lines) e.g. here
• [ ] anything else that might make code difficult to read or use?