From eb3081f1fab92c9c86fd90b0dbf26ec7f338b738 Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@jetbrains.com>
Date: Wed, 10 Apr 2013 20:49:26 +0400
Subject: [PATCH] Refactored collecting instance attributes

---
 .../python/psi/impl/PyClassImpl.java          | 44 ++++++++++---------
 1 file changed, 23 insertions(+), 21 deletions(-)

diff --git a/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java b/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java
index 2e0d9f5d5582..6dd9750154da 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java
@@ -935,32 +935,34 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
       return;
     }
     final PyFunctionStub methodStub = method.getStub();
-    if (methodStub != null) {
-      final PyTargetExpression[] targets = methodStub.getChildrenByType(PyElementTypes.TARGET_EXPRESSION, PyTargetExpression.EMPTY_ARRAY);
-      for (PyTargetExpression target : targets) {
-        if (!result.containsKey(target.getName())) {
-          result.put(target.getName(), target);
-        }
-      }
-    }
-    else {
-      final PyStatementList statementList = method.getStatementList();
-      if (statementList != null) {
-        statementList.accept(new PyRecursiveElementVisitor() {
-          public void visitPyAssignmentStatement(final PyAssignmentStatement node) {
-            collectNewTargets(result, node);
-          }
-        });
+    for (PyTargetExpression target : getTargetExpressions(method)) {
+      if (methodStub != null || PyUtil.isInstanceAttribute(target)) {
+        result.put(target.getName(), target);
       }
     }
   }
 
-  private static void collectNewTargets(Map<String, PyTargetExpression> collected, PyAssignmentStatement node) {
-    final PyExpression[] targets = node.getTargets();
-    for (PyExpression target : targets) {
-      if (target instanceof PyTargetExpression && PyUtil.isInstanceAttribute(target)) {
-        collected.put(target.getName(), (PyTargetExpression)target);
+  @NotNull
+  private static List<PyTargetExpression> getTargetExpressions(@NotNull PyFunction function) {
+    final PyFunctionStub stub = function.getStub();
+    if (stub != null) {
+      return Arrays.asList(stub.getChildrenByType(PyElementTypes.TARGET_EXPRESSION, PyTargetExpression.EMPTY_ARRAY));
+    }
+    else {
+      final PyStatementList statementList = function.getStatementList();
+      final List<PyTargetExpression> result = new ArrayList<PyTargetExpression>();
+      if (statementList != null) {
+        statementList.accept(new PyRecursiveElementVisitor() {
+          public void visitPyAssignmentStatement(final PyAssignmentStatement node) {
+            for (PyExpression expression : node.getTargets()) {
+              if (expression instanceof PyTargetExpression) {
+                result.add((PyTargetExpression)expression);
+              }
+            }
+          }
+        });
       }
+      return result;
     }
   }