diff options
-rw-r--r-- | ChangeLog | 2 | ||||
-rw-r--r-- | astroid/brain/brain_namedtuple_enum.py | 23 | ||||
-rw-r--r-- | astroid/tests/unittest_brain.py | 32 | ||||
-rw-r--r-- | astroid/tests/unittest_object_model.py | 16 |
4 files changed, 65 insertions, 8 deletions
@@ -2,6 +2,8 @@ Change log for the astroid package (used to be astng) ===================================================== -- + * `namedtuple` inference now understands `rename` keyword argument + * Classes can now know their definition-time arguments. Classes can support keyword arguments, which are passed when diff --git a/astroid/brain/brain_namedtuple_enum.py b/astroid/brain/brain_namedtuple_enum.py index d7c83bb8..b99cb761 100644 --- a/astroid/brain/brain_namedtuple_enum.py +++ b/astroid/brain/brain_namedtuple_enum.py @@ -8,11 +8,13 @@ import functools import sys +import keyword from textwrap import dedent from astroid import ( MANAGER, UseInferenceDefault, inference_tip, InferenceError) +from astroid import arguments from astroid import exceptions from astroid import nodes from astroid.builder import AstroidBuilder, extract_node @@ -113,6 +115,16 @@ def infer_named_tuple(node, context=None): """Specific inference function for namedtuple Call node""" class_node, name, attributes = infer_func_form(node, nodes.Tuple._proxied, context=context) + call_site = arguments.CallSite.from_call(node) + func = next(extract_node('import collections; collections.namedtuple').infer()) + try: + rename = next(call_site.infer_argument(func, 'rename', context)).bool_value() + except InferenceError: + rename = False + + if rename: + attributes = _get_renamed_namedtuple_atributes(attributes) + field_def = (" {name} = property(lambda self: self[{index:d}], " "doc='Alias for field number {index:d}')") field_defs = '\n'.join(field_def.format(name=name, index=index) @@ -142,6 +154,17 @@ class %(name)s(tuple): return iter([class_node]) +def _get_renamed_namedtuple_atributes(field_names): + names = list(field_names) + seen = set() + for i, name in enumerate(field_names): + if (not all(c.isalnum() or c == '_' for c in name) or keyword.iskeyword(name) + or not name or name[0].isdigit() or name.startswith('_') or name in seen): + names[i] = '_%d' % i + seen.add(name) + return tuple(names) + + def infer_enum(node, context=None): """ Specific inference function for enum Call node. """ enum_meta = extract_node(''' diff --git a/astroid/tests/unittest_brain.py b/astroid/tests/unittest_brain.py index 7fb0e94f..c85a61f6 100644 --- a/astroid/tests/unittest_brain.py +++ b/astroid/tests/unittest_brain.py @@ -161,6 +161,38 @@ class NamedTupleTest(unittest.TestCase): self.assertIn('field', inferred.locals) self.assertIn('other', inferred.locals) + def test_namedtuple_rename_keywords(self): + node = builder.extract_node(""" + from collections import namedtuple + Tuple = namedtuple("Tuple", "abc def", rename=True) + Tuple #@ + """) + inferred = next(node.infer()) + self.assertIn('abc', inferred.locals) + self.assertIn('_1', inferred.locals) + + def test_namedtuple_rename_duplicates(self): + node = builder.extract_node(""" + from collections import namedtuple + Tuple = namedtuple("Tuple", "abc abc abc", rename=True) + Tuple #@ + """) + inferred = next(node.infer()) + self.assertIn('abc', inferred.locals) + self.assertIn('_1', inferred.locals) + self.assertIn('_2', inferred.locals) + + def test_namedtuple_rename_uninferable(self): + node = builder.extract_node(""" + from collections import namedtuple + Tuple = namedtuple("Tuple", "a b c", rename=UNINFERABLE) + Tuple #@ + """) + inferred = next(node.infer()) + self.assertIn('a', inferred.locals) + self.assertIn('b', inferred.locals) + self.assertIn('c', inferred.locals) + class DefaultDictTest(unittest.TestCase): diff --git a/astroid/tests/unittest_object_model.py b/astroid/tests/unittest_object_model.py index fbe88159..0649edb5 100644 --- a/astroid/tests/unittest_object_model.py +++ b/astroid/tests/unittest_object_model.py @@ -31,7 +31,7 @@ class InstanceModelTest(unittest.TestCase): a.__module__ #@ a.__doc__ #@ a.__dict__ #@ - ''', module_name='collections') + ''', module_name='fake_module') cls = next(ast_nodes[0].infer()) self.assertIsInstance(cls, astroid.ClassDef) @@ -39,7 +39,7 @@ class InstanceModelTest(unittest.TestCase): module = next(ast_nodes[1].infer()) self.assertIsInstance(module, astroid.Const) - self.assertEqual(module.value, 'collections') + self.assertEqual(module.value, 'fake_module') doc = next(ast_nodes[2].infer()) self.assertIsInstance(doc, astroid.Const) @@ -197,11 +197,11 @@ class ClassModelTest(unittest.TestCase): A.__class__ #@ A.__dict__ #@ A.__subclasses__() #@ - ''', module_name='collections') + ''', module_name='fake_module') module = next(ast_nodes[0].infer()) self.assertIsInstance(module, astroid.Const) - self.assertEqual(module.value, 'collections') + self.assertEqual(module.value, 'fake_module') name = next(ast_nodes[1].infer()) self.assertIsInstance(name, astroid.Const) @@ -209,7 +209,7 @@ class ClassModelTest(unittest.TestCase): qualname = next(ast_nodes[2].infer()) self.assertIsInstance(qualname, astroid.Const) - self.assertEqual(qualname.value, 'collections.A') + self.assertEqual(qualname.value, 'fake_module.A') doc = next(ast_nodes[3].infer()) self.assertIsInstance(doc, astroid.Const) @@ -358,7 +358,7 @@ class FunctionModelTest(unittest.TestCase): func.__globals__ #@ func.__code__ #@ func.__closure__ #@ - ''', module_name='collections') + ''', module_name='fake_module') name = next(ast_nodes[0].infer()) self.assertIsInstance(name, astroid.Const) @@ -370,11 +370,11 @@ class FunctionModelTest(unittest.TestCase): qualname = next(ast_nodes[2].infer()) self.assertIsInstance(qualname, astroid.Const) - self.assertEqual(qualname.value, 'collections.func') + self.assertEqual(qualname.value, 'fake_module.func') module = next(ast_nodes[3].infer()) self.assertIsInstance(module, astroid.Const) - self.assertEqual(module.value, 'collections') + self.assertEqual(module.value, 'fake_module') defaults = next(ast_nodes[4].infer()) self.assertIsInstance(defaults, astroid.Tuple) |