diff options
Diffstat (limited to 'astroid/helpers.py')
-rw-r--r-- | astroid/helpers.py | 24 |
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
|