[minor] Added some type annotations