Add support for indexing containers with instances which provides an __index__ returning-int method.

This patch moves _class_as_index to helpers, where it becames class_instance_as_index.
Also, it instantiates its own call context, which makes certain idioms with lambdas
to work.

1 files changed, 24 insertions, 0 deletions

diff --git a/astroid/helpers.py b/astroid/helpers.py
index 2b318410..feb0985e 100644
--- a/astroid/helpers.py
+++ b/astroid/helpers.py
@@ -26,6 +26,7 @@ from astroid import bases
 from astroid import context as contextmod

 from astroid import exceptions

 from astroid import manager

+from astroid import nodes

 from astroid import raw_building

 from astroid import scoped_nodes

 from astroid import util

@@ -156,3 +157,26 @@ def is_subtype(type1, type2):
 def is_supertype(type1, type2):

     """Check if *type2* is a supertype of *type1*."""

     return _type_check(type1, type2)

+

+

+def class_instance_as_index(node):

+    """Get the value as an index for the given instance.

+

+    If an instance provides an __index__ method, then it can

+    be used in some scenarios where an integer is expected,

+    for instance when multiplying or subscripting a list.

+    """

+    context = contextmod.InferenceContext()

+    context.callcontext = contextmod.CallContext(args=[node])

+

+    try:

+        for inferred in node.igetattr('__index__', context=context):

+            if not isinstance(inferred, bases.BoundMethod):

+                continue

+

+            for result in inferred.infer_call_result(node, context=context):

+                if (isinstance(result, nodes.Const)

+                        and isinstance(result.value, int)):

+                    return result

+    except exceptions.InferenceError:

+        pass