Fix the __init/__exit annotations