From 109e28f53705a3de37946712ec9f8b7dd239d713 Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Thu, 8 Aug 2013 14:43:28 +0200
Subject: [PATCH] understand dict.update() calls in evaluators

---
 .../src/com/jetbrains/python/PyNames.java     |  1 +
 .../python/psi/impl/PyFileEvaluator.java      | 23 +++++++++++++++----
 .../jetbrains/python/PyFileEvaluatorTest.java |  7 ++++++
 3 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/python/psi-api/src/com/jetbrains/python/PyNames.java b/python/psi-api/src/com/jetbrains/python/PyNames.java
index 3143389c818e..da970ed07576 100644
--- a/python/psi-api/src/com/jetbrains/python/PyNames.java
+++ b/python/psi-api/src/com/jetbrains/python/PyNames.java
@@ -125,6 +125,7 @@ public class PyNames {
 
   public static final String KEYS = "keys";
   public static final String EXTEND = "extend";
+  public static final String UPDATE = "update";
 
   public static final String PASS = "pass";
 
diff --git a/python/src/com/jetbrains/python/psi/impl/PyFileEvaluator.java b/python/src/com/jetbrains/python/psi/impl/PyFileEvaluator.java
index 1ca8b5407cd3..a33a41616a2f 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyFileEvaluator.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyFileEvaluator.java
@@ -96,13 +96,16 @@ public class PyFileEvaluator {
         PyExpression callee = node.getCallee();
         if (callee instanceof PyReferenceExpression) {
           PyReferenceExpression calleeRef = (PyReferenceExpression)callee;
-          if (PyNames.EXTEND.equals(calleeRef.getReferencedName()) && node.getArguments().length == 1) {
-            PyExpression qualifier = calleeRef.getQualifier();
-            if (qualifier instanceof PyReferenceExpression) {
-              PyReferenceExpression qualifierRef = (PyReferenceExpression)qualifier;
-              if (qualifierRef.getQualifier() == null) {
+          PyExpression qualifier = calleeRef.getQualifier();
+          if (qualifier instanceof PyReferenceExpression) {
+            PyReferenceExpression qualifierRef = (PyReferenceExpression)qualifier;
+            if (qualifierRef.getQualifier() == null) {
+              if (PyNames.EXTEND.equals(calleeRef.getReferencedName()) && node.getArguments().length == 1) {
                 processExtendCall(node, qualifierRef.getReferencedName());
               }
+              else if (PyNames.UPDATE.equals(calleeRef.getReferencedName()) && node.getArguments().length == 1) {
+                processUpdateCall(node, qualifierRef.getReferencedName());
+              }
             }
           }
         }
@@ -151,6 +154,16 @@ public class PyFileEvaluator {
     }
   }
 
+  private void processUpdateCall(PyCallExpression node, String name) {
+    Object value = myNamespace.get(name);
+    if (value instanceof Map) {
+      Object argValue = createEvaluator().evaluate(node.getArguments()[0]);
+      if (argValue instanceof Map) {
+        ((Map)value).putAll((Map)argValue);
+      }
+    }
+  }
+
   private PyEvaluator createEvaluator() {
     return new PyPathEvaluator(myCurrentFilePath);
   }
diff --git a/python/testSrc/com/jetbrains/python/PyFileEvaluatorTest.java b/python/testSrc/com/jetbrains/python/PyFileEvaluatorTest.java
index d689235de9f4..5b1e51b1c7a5 100644
--- a/python/testSrc/com/jetbrains/python/PyFileEvaluatorTest.java
+++ b/python/testSrc/com/jetbrains/python/PyFileEvaluatorTest.java
@@ -59,6 +59,13 @@ public class PyFileEvaluatorTest extends PyTestCase {
     assertEquals("c", map.get("b"));
   }
 
+  public void testDictUpdate() {
+    PyFileEvaluator eval = doEvaluate("a={}\na.update({'b': 'c'})");
+    Map map = (Map) eval.getValue("a");
+    assertEquals(1, map.size());
+    assertEquals("c", map.get("b"));
+  }
+
   private PyFileEvaluator doEvaluate(String text) {
     PyFileEvaluator eval = new PyFileEvaluator();
     PyFile file = (PyFile)PsiFileFactory.getInstance(myFixture.getProject()).createFileFromText("a.py", PythonFileType.INSTANCE, text);