Fix the `with` context
Add a missed `with`
Fix imports