More type annotation updates