Improve `add_input` and `add_output` methods

Redesign, simplify and fix the type annotations