Search code examples
pythonpycharmgeneratordocstring

PyCharm: How to document :rtype: for function that returns generator


I try to document a :rtype: docstring parameter returning a generator of Node:

def __iter__(self):
    """iterate over node children

    :rtype: ???
    """
    for node in self.children.itervalues():
        yield node

What :rtype: is supposed to be? generator of Node doesn't seems to work.


Solution

  • You can hint the type of self.children using :type below it, as described in the official documentation. This would let you work with this variable everywhere using defined types. Here is an example:

    class Node(object):
        def foo(self):
            pass
    
    
    class Iterator(object):
        def __init__(self):
            self.children = ''
            """:type: dict[str, Node]"""   # <<< define type of self.children after assigning the value
    
        def __iter__(self):
            for node in self.children.itervalues():
                yield node
    
    
    if __name__ == '__main__':
        for node in Iterator():
            node. # <<< type hinting will work here
    

    enter image description here

    But if you prefer to set the return type of __iter__ method, you can do it using next syntax:

    def __iter__(self):
        """
        :rtype: collections.Iterable[Node]
        """
        for node in self.children.itervalues():
            yield node
    

    I would prefer first solution though.