Add support for inference of subscript operations on dict literals. Closes #123074

actually fix Dict.getitem and enhance infer_subscript

--HG--
branch : stable

4 files changed, 33 insertions, 16 deletions

diff --git a/ChangeLog b/ChangeLog
index c8dafdc3..79e931a7 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -10,6 +10,8 @@ Change log for the astng package
     * #123068: Make sure .as_string() returns valid code for yields in
       expressions.
     * #47957: Set literals are now correctly treated as inference leaves.
+    * #123074: Add support for inference of subscript operations on dict 
+      literals.
 
 2013-02-27  -- 0.24.2
     * pylint-brain: more subprocess.Popen faking (see #46273)
diff --git a/inference.py b/inference.py
index 49d2424a..e2d70d4d 100644
--- a/inference.py
+++ b/inference.py
@@ -1,4 +1,4 @@
-# copyright 2003-2011 LOGILAB S.A. (Paris, FRANCE), all rights reserved.
+# copyright 2003-2013 LOGILAB S.A. (Paris, FRANCE), all rights reserved.
 # contact http://www.logilab.fr/ -- mailto:contact@logilab.fr
 # copyright 2003-2010 Sylvain Thenault, all rights reserved.
 # contact mailto:thenault@gmail.com
@@ -238,15 +238,19 @@ nodes.Global.infer = path_wrapper(infer_global)
 
 def infer_subscript(self, context=None):
     """infer simple subscription such as [1,2,3][0] or (1,2,3)[-1]"""
-    if isinstance(self.slice, nodes.Index):
-        index = self.slice.value.infer(context).next()
-        if index is YES:
-            yield YES
-            return
+    value = self.value.infer(context).next()
+    if value is YES:
+        yield YES
+        return
+
+    index = self.slice.infer(context).next()
+    if index is YES:
+        yield YES
+        return
+
+    if isinstance(index, nodes.Const):
         try:
-            # suppose it's a Tuple/List node (attribute error else)
-            # XXX infer self.value?
-            assigned = self.value.getitem(index.value, context)
+            assigned = value.getitem(index.value, context)
         except AttributeError:
             raise InferenceError()
         except (IndexError, TypeError):
@@ -381,3 +385,7 @@ def infer_empty_node(self, context=None):
             yield YES
 nodes.EmptyNode.infer = path_wrapper(infer_empty_node)
 
+
+def infer_index(self, context=None):
+    return self.value.infer(context)
+nodes.Index.infer = infer_index
diff --git a/node_classes.py b/node_classes.py
index 3073535e..e3f1c8db 100644
--- a/node_classes.py
+++ b/node_classes.py
@@ -1,4 +1,4 @@
-# copyright 2003-2012 LOGILAB S.A. (Paris, FRANCE), all rights reserved.
+# copyright 2003-2013 LOGILAB S.A. (Paris, FRANCE), all rights reserved.
 # contact http://www.logilab.fr/ -- mailto:contact@logilab.fr
 # copyright 2003-2010 Sylvain Thenault, all rights reserved.
 # contact mailto:thenault@gmail.com
@@ -513,14 +513,16 @@ class Dict(NodeNG, Instance):
     def itered(self):
         return self.items[::2]
 
-    def getitem(self, key, context=None):
-        for i in xrange(0, len(self.items), 2):
-            for inferedkey in self.items[i].infer(context):
+    def getitem(self, lookup_key, context=None):
+        for key, value in self.items:
+            for inferedkey in key.infer(context):
                 if inferedkey is YES:
                     continue
-                if isinstance(inferedkey, Const) and inferedkey.value == key:
-                    return self.items[i+1]
-        raise IndexError(key)
+                if isinstance(inferedkey, Const) and inferedkey.value == lookup_key:
+                    return value
+        # This should raise KeyError, but all call sites only catch
+        # IndexError. Let's leave it like that for now.
+        raise IndexError(lookup_key)
 
 
 class Discard(Statement):
diff --git a/test/unittest_inference.py b/test/unittest_inference.py
index 3853e353..438ee265 100644
--- a/test/unittest_inference.py
+++ b/test/unittest_inference.py
@@ -627,6 +627,9 @@ b = (1, 2, 3)[1]
 c = (1, 2, 3)[-1]
 d = a + b + c
 print (d)
+e = {'key': 'value'}
+f = e['key']
+print (f)
         '''
         astng = builder.string_build(code, __name__, __file__)
         self.assertEqual([i.value for i in
@@ -637,6 +640,8 @@ print (d)
                                 get_name_node(astng, 'c', -1).infer()], [3])
         self.assertEqual([i.value for i in
                                 get_name_node(astng, 'd', -1).infer()], [6])
+        self.assertEqual([i.value for i in
+                          get_name_node(astng, 'f', -1).infer()], ['value'])
 
     #def test_simple_tuple(self):
         #"""test case for a simple tuple value"""