Add mypy compliance

## Description

Opening this issue to address a PR already opened by @sdavtaker . 
It would be nice to have type hints to make mypy validation pass.

## Proposed solution

#111 is already opened