From 59c975f00e1838c2116690f7966c0d348c49033e Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@gmail.com>
Date: Thu, 16 Aug 2012 18:55:50 +0400
Subject: [PATCH] Refactored searching for fake '__generator' class

---
 python/psi-api/src/com/jetbrains/python/PyNames.java        | 2 ++
 .../src/com/jetbrains/python/psi/impl/PyFunctionImpl.java   | 6 +++---
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/python/psi-api/src/com/jetbrains/python/PyNames.java b/python/psi-api/src/com/jetbrains/python/PyNames.java
index 76466c996440..e256c462a2f2 100644
--- a/python/psi-api/src/com/jetbrains/python/PyNames.java
+++ b/python/psi-api/src/com/jetbrains/python/PyNames.java
@@ -37,7 +37,9 @@ public class PyNames {
   public static final String NONE = "None";
   public static final String TRUE = "True";
   public static final String FALSE = "False";
+
   public static final String FAKE_OLD_BASE = "___Classobj";
+  public static final String FAKE_GENERATOR = "__generator";
 
   public static final String FUTURE_MODULE = "__future__";
   public static final String UNICODE_LITERALS = "unicode_literals";
diff --git a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
index 7168998b8506..69e36e05e9c0 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
@@ -270,9 +270,9 @@ public class PyFunctionImpl extends PyPresentableElementImpl<PyFunctionStub> imp
       }
     }
     if (elementType != null) {
-      final PyType it = PyTypeParser.getTypeByName(this, "__generator");
-      if (it instanceof PyClassType) {
-        return new PyCollectionTypeImpl(((PyClassType)it).getPyClass(), false, elementType.get());
+      final PyClass generator = cache.getClass(PyNames.FAKE_GENERATOR);
+      if (generator != null) {
+        return new PyCollectionTypeImpl(generator, false, elementType.get());
       }
     }
     return null;