PY-26032 data and target attributes are unresolved in scikit Bunch class

This commit is contained in:
Ekaterina Tuzova
2017-09-11 15:25:30 +03:00
parent 04bf66e450
commit cbcb17f65d
@@ -15,10 +15,12 @@
*/
package com.jetbrains.numpy.codeInsight;
import com.google.common.collect.Lists;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiReference;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.codeInsight.PyCustomMember;
import com.jetbrains.python.psi.PyClass;
import com.jetbrains.python.psi.PyFunction;
import com.jetbrains.python.psi.types.PyClassMembersProviderBase;
import com.jetbrains.python.psi.types.PyClassType;
@@ -35,25 +37,36 @@ import java.util.List;
* User : ktisha
*/
public class NumpyClassMembersProvider extends PyClassMembersProviderBase {
private static final String BUNCH = "sklearn.datasets.base.Bunch";
public static final List<String> BUNCH_MEMBERS = Lists.newArrayList("target", "data", "filenames", "target_names", "DESCR");
@NotNull
@Override
public Collection<PyCustomMember> getMembers(PyClassType clazz, PsiElement location, TypeEvalContext typeEvalContext) {
if (location != null && clazz.getPyClass().isSubclass(PyNames.TYPES_FUNCTION_TYPE, typeEvalContext)) {
final PsiElement element = location.getOriginalElement();
final PsiReference reference = element.getReference();
if (reference != null) {
final PsiElement resolved = reference.resolve();
if (resolved instanceof PyFunction) {
final List<PyCustomMember> result = new ArrayList<>();
if (NumpyUfuncs.isUFunc(((PyFunction)resolved).getName()) && NumpyDocStringTypeProvider.isInsideNumPy(resolved)) {
for (String method : NumpyUfuncs.UFUNC_METHODS) {
result.add(new PyCustomMember(method, resolved));
if (location != null) {
final PyClass pyClass = clazz.getPyClass();
if (pyClass.isSubclass(PyNames.TYPES_FUNCTION_TYPE, typeEvalContext)) {
final PsiElement element = location.getOriginalElement();
final PsiReference reference = element.getReference();
if (reference != null) {
final PsiElement resolved = reference.resolve();
if (resolved instanceof PyFunction) {
final List<PyCustomMember> result = new ArrayList<>();
if (NumpyUfuncs.isUFunc(((PyFunction)resolved).getName()) && NumpyDocStringTypeProvider.isInsideNumPy(resolved)) {
for (String method : NumpyUfuncs.UFUNC_METHODS) {
result.add(new PyCustomMember(method, resolved));
}
return result;
}
return result;
}
}
}
else if (BUNCH.equals(pyClass.getQualifiedName())) {
final List<PyCustomMember> result = new ArrayList<>();
for (String member : BUNCH_MEMBERS) {
result.add(new PyCustomMember(member, NumpyDocStringTypeProvider.NDARRAY, true));
}
return result;
}
}
return Collections.emptyList();
}