diff --git a/python/src/com/jetbrains/numpy/codeInsight/NumpyClassMembersProvider.java b/python/src/com/jetbrains/numpy/codeInsight/NumpyClassMembersProvider.java index 37efa9088246..58499183a6e4 100644 --- a/python/src/com/jetbrains/numpy/codeInsight/NumpyClassMembersProvider.java +++ b/python/src/com/jetbrains/numpy/codeInsight/NumpyClassMembersProvider.java @@ -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 BUNCH_MEMBERS = Lists.newArrayList("target", "data", "filenames", "target_names", "DESCR"); @NotNull @Override public Collection 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 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 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 result = new ArrayList<>(); + for (String member : BUNCH_MEMBERS) { + result.add(new PyCustomMember(member, NumpyDocStringTypeProvider.NDARRAY, true)); + } + return result; + } } return Collections.emptyList(); }