PyCharm doesn't recognize iterator and enumerator return types

Question

Consider the following example of a very simple iterator:

class A(object):
    def __init__(self, n):
        self.list = [0] * n

    def __iter__(self):
        return AIter(self)

and

class AIter(object):
    def __init__(self, a):
        """
        Args:
            a(A):
        """
        self.obj = a
        self.n = len(a.list)
        self.i = 0

    def __next__(self):
        if self.i == self.n:
            raise StopIteration
        else:
            self.i += 1
            return self.obj.list[self.i-1]

When typing in a PyCharm editor, after defining a=A(5), a.list[0]. does complete the functions for int. However, when using something like for j in a: I cannot get PyCharm to understand that j is an int. It knows that AIter.next() returns an int, and it knows that A.__iter__() returns an AIter.

This can be solved by writing # type: int after the for, but this seems like a workaround, and I have to remember to do this every time I call a for or enumerate. Can I make PyCharm recognize the correct types automatically?

Answer 1

Fixed since February 15, 2017. The latest EAP is available here.