PyCharm getitem warning for functions with arrays

Question

I'm getting code inspection warnings from PyCharm. I understand the logic, but I'm not clear on the appropriate way to fix it. Say I have the following example function:

def get_ydata(xdata):
    ydata = xdata ** 2
    for i in range(len(ydata)):
        print ydata[i]
return ydata

I get 2 warnings:

>> Expected type 'Sized', got 'int' instead (at line 3)
>> Class 'int' does not define '__getitem__', so the '[]' operator cannot be used on its instances (at line 4)

The purpose of the function is of course to parse a numpy array of xdata. But PyCharm doesn't know that, so without any further indication assumes that xdata (and therefore also ydata) is an integer.

What is the appropriate way to address this warning? I should note that adding a type checking line will fix the warning. Is that the optimal solution? For example:

if not type(ydata) is np.ndarray:
    ydata = np.array(ydata)

Lastly, adding Sphinx docstring information does not seem to have any effect on the warnings. (warning still sees 'int' when xdata is specified as str). Also iterating over y directly results in the following error:

for y in ydata:
...
>> Expected 'collections.Iterable', got 'int' instead

Solution

Pycharm has type hinting features that may be of use.

For example in this case, the following code makes the errors go away:

import numpy as np

def get_ydata(xdata):
    ydata = xdata ** 2  # type: np.ndarray
    for i in range(len(ydata)):
        print(ydata[i])
    return ydata

Recent python versions also include support for type annotations

import numpy as np
def get_ydata(xdata: np.ndarray):
    ...