diff options
| -rw-r--r-- | ChangeLog | 2 | ||||
| -rw-r--r-- | brain/py2stdlib.py | 22 | ||||
| -rw-r--r-- | test/unittest_brain.py | 25 |
3 files changed, 44 insertions, 5 deletions
@@ -2,6 +2,8 @@ Change log for the astroid package (used to be astng) ===================================================== -- + * Add support for inferring arguments to namedtuple invocations. + * Make sure that objects returned for namedtuple inference have parents. diff --git a/brain/py2stdlib.py b/brain/py2stdlib.py index 31a2477e..27507ec4 100644 --- a/brain/py2stdlib.py +++ b/brain/py2stdlib.py @@ -5,7 +5,8 @@ Currently help understanding of : * hashlib.md5 and hashlib.sha1 """ -from astroid import MANAGER, AsStringRegexpPredicate, UseInferenceDefault, inference_tip +from astroid import MANAGER, AsStringRegexpPredicate, UseInferenceDefault, inference_tip, YES +from astroid import exceptions from astroid import nodes from astroid.builder import AstroidBuilder @@ -183,6 +184,16 @@ MODULE_TRANSFORMS['subprocess'] = subprocess_transform def infer_named_tuple(node, context=None): """Specific inference function for namedtuple CallFunc node""" + def infer_first(node): + try: + value = node.infer().next() + if value is YES: + raise UseInferenceDefault() + else: + return value + except StopIteration: + raise InferenceError() + # node is a CallFunc node, class name as first argument and generated class # attributes as second argument if len(node.args) != 2: @@ -191,12 +202,13 @@ def infer_named_tuple(node, context=None): # namedtuple list of attributes can be a list of strings or a # whitespace-separate string try: - name = node.args[0].value + name = infer_first(node.args[0]).value + names = infer_first(node.args[1]) try: - attributes = node.args[1].value.split() + attributes = names.value.split() except AttributeError: - attributes = [const.value for const in node.args[1].elts] - except AttributeError: + attributes = [infer_first(const).value for const in names.elts] + except (AttributeError, exceptions.InferenceError): raise UseInferenceDefault() # we want to return a Class node instance with proper attributes set class_node = nodes.Class(name, 'docstring') diff --git a/test/unittest_brain.py b/test/unittest_brain.py index dac25e62..7b750cf6 100644 --- a/test/unittest_brain.py +++ b/test/unittest_brain.py @@ -19,6 +19,7 @@ from logilab.common.testlib import TestCase, unittest_main from astroid import MANAGER +from astroid import bases from astroid import test_utils import astroid @@ -52,5 +53,29 @@ class NamedTupleTest(TestCase): for anc in klass.ancestors(): self.assertFalse(anc.parent is None) + def test_namedtuple_inference(self): + klass = test_utils.extract_node(""" + from collections import namedtuple + + name = "X" + fields = ["a", "b", "c"] + class X(namedtuple(name, fields)): + pass + """) + for base in klass.ancestors(): + if base.name == 'X': + break + self.assertItemsEqual(["a", "b", "c"], base.instance_attrs.keys()) + + def test_namedtuple_inference_failure(self): + klass = test_utils.extract_node(""" + from collections import namedtuple + + def foo(fields): + return __(namedtuple("foo", fields)) + """) + self.assertIs(bases.YES, klass.infer().next()) + + if __name__ == '__main__': unittest_main() |