From 90d83aab32afe1dc9ac9651efb1994b17d22ca37 Mon Sep 17 00:00:00 2001 From: Claudiu Popa Date: Thu, 22 Oct 2015 21:40:12 +0100 Subject: 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. --- astroid/helpers.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) (limited to 'astroid/helpers.py') 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 -- cgit v1.2.1