diff options
Diffstat (limited to 'astroid/brain/brain_namedtuple_enum.py')
| -rw-r--r-- | astroid/brain/brain_namedtuple_enum.py | 209 |
1 files changed, 105 insertions, 104 deletions
diff --git a/astroid/brain/brain_namedtuple_enum.py b/astroid/brain/brain_namedtuple_enum.py index cc399246..85390a82 100644 --- a/astroid/brain/brain_namedtuple_enum.py +++ b/astroid/brain/brain_namedtuple_enum.py @@ -21,9 +21,7 @@ import functools import keyword from textwrap import dedent -from astroid import ( - MANAGER, UseInferenceDefault, inference_tip, - InferenceError) +from astroid import MANAGER, UseInferenceDefault, inference_tip, InferenceError from astroid import arguments from astroid import exceptions from astroid import nodes @@ -31,11 +29,8 @@ from astroid.builder import AstroidBuilder, extract_node from astroid import util -TYPING_NAMEDTUPLE_BASENAMES = { - 'NamedTuple', - 'typing.NamedTuple' -} -ENUM_BASE_NAMES = {'Enum', 'IntEnum', 'enum.Enum', 'enum.IntEnum'} +TYPING_NAMEDTUPLE_BASENAMES = {"NamedTuple", "typing.NamedTuple"} +ENUM_BASE_NAMES = {"Enum", "IntEnum", "enum.Enum", "enum.IntEnum"} def _infer_first(node, context): @@ -52,7 +47,6 @@ def _infer_first(node, context): def _find_func_form_arguments(node, context): - def _extract_namedtuple_arg_or_keyword(position, key_name=None): if len(args) > position: @@ -62,18 +56,12 @@ def _find_func_form_arguments(node, context): args = node.args keywords = node.keywords - found_keywords = { - keyword.arg: keyword.value for keyword in keywords - } if keywords else {} - - name = _extract_namedtuple_arg_or_keyword( - position=0, - key_name='typename' - ) - names = _extract_namedtuple_arg_or_keyword( - position=1, - key_name='field_names' + found_keywords = ( + {keyword.arg: keyword.value for keyword in keywords} if keywords else {} ) + + name = _extract_namedtuple_arg_or_keyword(position=0, key_name="typename") + names = _extract_namedtuple_arg_or_keyword(position=1, key_name="field_names") if name and names: return name.value, names @@ -90,30 +78,35 @@ def infer_func_form(node, base_type, context=None, enum=False): try: name, names = _find_func_form_arguments(node, context) try: - attributes = names.value.replace(',', ' ').split() + attributes = names.value.replace(",", " ").split() except AttributeError: if not enum: - attributes = [_infer_first(const, context).value - for const in names.elts] + attributes = [ + _infer_first(const, context).value for const in names.elts + ] else: # Enums supports either iterator of (name, value) pairs # or mappings. - if hasattr(names, 'items') and isinstance(names.items, list): - attributes = [_infer_first(const[0], context).value - for const in names.items - if isinstance(const[0], nodes.Const)] - elif hasattr(names, 'elts'): + if hasattr(names, "items") and isinstance(names.items, list): + attributes = [ + _infer_first(const[0], context).value + for const in names.items + if isinstance(const[0], nodes.Const) + ] + elif hasattr(names, "elts"): # Enums can support either ["a", "b", "c"] # or [("a", 1), ("b", 2), ...], but they can't # be mixed. - if all(isinstance(const, nodes.Tuple) - for const in names.elts): - attributes = [_infer_first(const.elts[0], context).value - for const in names.elts - if isinstance(const, nodes.Tuple)] + if all(isinstance(const, nodes.Tuple) for const in names.elts): + attributes = [ + _infer_first(const.elts[0], context).value + for const in names.elts + if isinstance(const, nodes.Tuple) + ] else: - attributes = [_infer_first(const, context).value - for const in names.elts] + attributes = [ + _infer_first(const, context).value for const in names.elts + ] else: raise AttributeError if not attributes: @@ -123,9 +116,9 @@ def infer_func_form(node, base_type, context=None, enum=False): # If we can't infer the name of the class, don't crash, up to this point # we know it is a namedtuple anyway. - name = name or 'Uninferable' + name = name or "Uninferable" # we want to return a Class node instance with proper attributes set - class_node = nodes.ClassDef(name, 'docstring') + class_node = nodes.ClassDef(name, "docstring") class_node.parent = node.parent # set base class=tuple class_node.bases.append(base_type) @@ -156,39 +149,39 @@ def _looks_like(node, name): return False -_looks_like_namedtuple = functools.partial(_looks_like, name='namedtuple') -_looks_like_enum = functools.partial(_looks_like, name='Enum') -_looks_like_typing_namedtuple = functools.partial(_looks_like, name='NamedTuple') +_looks_like_namedtuple = functools.partial(_looks_like, name="namedtuple") +_looks_like_enum = functools.partial(_looks_like, name="Enum") +_looks_like_typing_namedtuple = functools.partial(_looks_like, name="NamedTuple") def infer_named_tuple(node, context=None): """Specific inference function for namedtuple Call node""" - tuple_base_name = nodes.Name(name='tuple', parent=node.root()) + tuple_base_name = nodes.Name(name="tuple", parent=node.root()) class_node, name, attributes = infer_func_form( - node, - tuple_base_name, - context=context, + node, tuple_base_name, context=context ) call_site = arguments.CallSite.from_call(node) - func = next(extract_node('import collections; collections.namedtuple').infer()) + func = next(extract_node("import collections; collections.namedtuple").infer()) try: - rename = next(call_site.infer_argument(func, 'rename', context)).bool_value() + rename = next(call_site.infer_argument(func, "rename", context)).bool_value() except InferenceError: rename = False if rename: attributes = _get_renamed_namedtuple_attributes(attributes) - replace_args = ', '.join( - '{arg}=None'.format(arg=arg) - for arg in attributes - ) + replace_args = ", ".join("{arg}=None".format(arg=arg) for arg in 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) - for index, name in enumerate(attributes)) - fake = AstroidBuilder(MANAGER).string_build(''' + 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) + for index, name in enumerate(attributes) + ) + fake = AstroidBuilder(MANAGER).string_build( + """ class %(name)s(tuple): __slots__ = () _fields = %(fields)r @@ -202,14 +195,18 @@ class %(name)s(tuple): def __getnewargs__(self): return tuple(self) %(field_defs)s - ''' % {'name': name, - 'fields': attributes, - 'field_defs': field_defs, - 'replace_args': replace_args}) - class_node.locals['_asdict'] = fake.body[0].locals['_asdict'] - class_node.locals['_make'] = fake.body[0].locals['_make'] - class_node.locals['_replace'] = fake.body[0].locals['_replace'] - class_node.locals['_fields'] = fake.body[0].locals['_fields'] + """ + % { + "name": name, + "fields": attributes, + "field_defs": field_defs, + "replace_args": replace_args, + } + ) + class_node.locals["_asdict"] = fake.body[0].locals["_asdict"] + class_node.locals["_make"] = fake.body[0].locals["_make"] + class_node.locals["_replace"] = fake.body[0].locals["_replace"] + class_node.locals["_fields"] = fake.body[0].locals["_fields"] for attr in attributes: class_node.locals[attr] = fake.body[0].locals[attr] # we use UseInferenceDefault, we can't be a generator so return an iterator @@ -220,16 +217,23 @@ def _get_renamed_namedtuple_attributes(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 + 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(''' + enum_meta = extract_node( + """ class EnumMeta(object): 'docstring' def __call__(self, node): @@ -255,9 +259,9 @@ def infer_enum(node, context=None): return Value() __members__ = [''] - ''') - class_node = infer_func_form(node, enum_meta, - context=context, enum=True)[0] + """ + ) + class_node = infer_func_form(node, enum_meta, context=context, enum=True)[0] return iter([class_node.instantiate_class()]) @@ -268,12 +272,11 @@ def infer_enum_class(node): # is a hack to support enums. if basename not in ENUM_BASE_NAMES: continue - if node.root().name == 'enum': + if node.root().name == "enum": # Skip if the class is directly from enum module. break for local, values in node.locals.items(): - if any(not isinstance(value, nodes.AssignName) - for value in values): + if any(not isinstance(value, nodes.AssignName) for value in values): continue targets = [] @@ -296,7 +299,8 @@ def infer_enum_class(node): new_targets = [] for target in targets: # Replace all the assignments with our mocked class. - classdef = dedent(''' + classdef = dedent( + """ class {name}({types}): @property def value(self): @@ -304,11 +308,12 @@ def infer_enum_class(node): @property def name(self): return {name} - '''.format( - name=target.name, - types=', '.join(node.basenames), - return_value=inferred_return_value, - )) + """.format( + name=target.name, + types=", ".join(node.basenames), + return_value=inferred_return_value, + ) + ) fake = AstroidBuilder(MANAGER).string_build(classdef)[target.name] fake.parent = target.parent for method in node.mymethods(): @@ -323,21 +328,21 @@ def infer_typing_namedtuple_class(class_node, context=None): """Infer a subclass of typing.NamedTuple""" # Check if it has the corresponding bases annassigns_fields = [ - annassign.target.name for annassign in class_node.body + annassign.target.name + for annassign in class_node.body if isinstance(annassign, nodes.AnnAssign) ] - code = dedent(''' + code = dedent( + """ from collections import namedtuple namedtuple({typename!r}, {fields!r}) - ''').format( - typename=class_node.name, - fields=",".join(annassigns_fields) - ) + """ + ).format(typename=class_node.name, fields=",".join(annassigns_fields)) node = extract_node(code) generated_class_node = next(infer_named_tuple(node, context)) for method in class_node.mymethods(): generated_class_node.locals[method.name] = [method] - return iter((generated_class_node, )) + return iter((generated_class_node,)) def infer_typing_namedtuple(node, context=None): @@ -349,7 +354,7 @@ def infer_typing_namedtuple(node, context=None): except InferenceError: raise UseInferenceDefault - if func.qname() != 'typing.NamedTuple': + if func.qname() != "typing.NamedTuple": raise UseInferenceDefault if len(node.args) != 2: @@ -367,31 +372,27 @@ def infer_typing_namedtuple(node, context=None): names.append(elt.elts[0].as_string()) typename = node.args[0].as_string() - node = extract_node('namedtuple(%(typename)s, (%(fields)s,)) ' % - {'typename': typename, 'fields': ",".join(names)}) + node = extract_node( + "namedtuple(%(typename)s, (%(fields)s,)) " + % {"typename": typename, "fields": ",".join(names)} + ) return infer_named_tuple(node, context) MANAGER.register_transform( - nodes.Call, inference_tip(infer_named_tuple), - _looks_like_namedtuple, + nodes.Call, inference_tip(infer_named_tuple), _looks_like_namedtuple ) +MANAGER.register_transform(nodes.Call, inference_tip(infer_enum), _looks_like_enum) MANAGER.register_transform( - nodes.Call, inference_tip(infer_enum), - _looks_like_enum, -) -MANAGER.register_transform( - nodes.ClassDef, infer_enum_class, - predicate=lambda cls: any(basename for basename in cls.basenames - if basename in ENUM_BASE_NAMES) + nodes.ClassDef, + infer_enum_class, + predicate=lambda cls: any( + basename for basename in cls.basenames if basename in ENUM_BASE_NAMES + ), ) MANAGER.register_transform( - nodes.ClassDef, - inference_tip(infer_typing_namedtuple_class), - _has_namedtuple_base, + nodes.ClassDef, inference_tip(infer_typing_namedtuple_class), _has_namedtuple_base ) MANAGER.register_transform( - nodes.Call, - inference_tip(infer_typing_namedtuple), - _looks_like_typing_namedtuple, + nodes.Call, inference_tip(infer_typing_namedtuple), _looks_like_typing_namedtuple ) |