consistent behavior of PyFunction.getContainingClass() in stub-based and non-stub-based cases (PY-1448)

python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java

@@ -23,6 +23,7 @@ import com.intellij.openapi.util.Ref;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.StubBasedPsiElement;
 import com.intellij.psi.search.LocalSearchScope;
 import com.intellij.psi.search.SearchScope;
 import com.intellij.psi.stubs.IStubElementType;
@@ -158,12 +159,9 @@ public class PyFunctionImpl extends PyPresentableElementImpl<PyFunctionStub> imp
       return null;
     }
 
-    final PsiElement parent = getParent();
-    if (parent instanceof PyStatementList) {
-      PsiElement pparent = parent.getParent();
-      if (pparent instanceof PyClass) {
-        return (PyClass)pparent;
-      }
+    final PsiElement parent = PsiTreeUtil.getParentOfType(this, StubBasedPsiElement.class);
+    if (parent instanceof PyClass) {
+      return (PyClass)parent;
     }
     return null;
   }

python/testData/findUsages/ConditionalFunctions.py

@@ -0,0 +1,14 @@
+import sys
+
+var = (sys.platform == 'win32')
+
+class A():
+    def __init__(self):
+        self.<caret>a = None
+
+    if var:
+        def func(self):
+            self.a = ""
+    else:
+        def func(self):
+            self.a = ()

python/testSrc/com/jetbrains/python/PyFindUsagesTest.java

@@ -147,6 +147,11 @@ public class PyFindUsagesTest extends PyTestCase {
     assertEquals(2, usages.size());
   }
 
+  public void testConditionalFunctions() {  // PY-1448
+    final Collection<UsageInfo> usages = myFixture.testFindUsages("findUsages/ConditionalFunctions.py");
+    assertEquals(3, usages.size());
+  }
+
   private Collection<UsageInfo> findMultiFileUsages(String filename) {
     final String testName = getTestName(false);
     myFixture.copyDirectoryToProject("findUsages/" + testName, "");