diff --git a/ChangeLog b/ChangeLog index 1d13c38408..c191dd2ca3 100644 --- a/ChangeLog +++ b/ChangeLog @@ -6,6 +6,9 @@ What's New in astroid 2.6.0? ============================ Release Date: TBA +* Adds ``ImportAlias`` node type included as children of ``Import`` and + ``ImportFrom`` nodes. + * Update enum brain to improve inference of .name and .value dynamic class attributes diff --git a/astroid/node_classes.py b/astroid/node_classes.py index 2e3e5f4f30..ae5dfdf9dc 100644 --- a/astroid/node_classes.py +++ b/astroid/node_classes.py @@ -3219,7 +3219,7 @@ def get_children(self): yield self.value -class ImportFrom(mixins.NoChildrenMixin, mixins.ImportFromMixin, Statement): +class ImportFrom(mixins.ImportFromMixin, Statement): """Class representing an :class:`ast.ImportFrom` node. >>> node = astroid.extract_node('from my_package import my_module') @@ -3227,11 +3227,23 @@ class ImportFrom(mixins.NoChildrenMixin, mixins.ImportFromMixin, Statement): """ + _astroid_fields = ("aliases",) + aliases = None + """ + :type: list(ImportAlias) or None + """ + _other_fields = ("modname", "names", "level") + names = None + """What is being imported from the module. - def __init__( - self, fromname, names, level=0, lineno=None, col_offset=None, parent=None - ): + Each entry is a :class:`tuple` of the name being imported, + and the alias that the name is assigned to (if any). + + :type: list(tuple(str, str or None)) + """ + + def __init__(self, fromname, level=0, lineno=None, col_offset=None, parent=None): """ :param fromname: The module that is being imported from. :type fromname: str or None @@ -3259,16 +3271,6 @@ def __init__( :type: str or None """ - - self.names = names - """What is being imported from the module. - - Each entry is a :class:`tuple` of the name being imported, - and the alias that the name is assigned to (if any). - - :type: list(tuple(str, str or None)) - """ - self.level = level """The level of relative import. @@ -3280,6 +3282,18 @@ def __init__( super().__init__(lineno, col_offset, parent) + def postinit(self, aliases=None): + """Do some setup after initialisation. + + :param aliases: The names being imported. + :type aliases: list(ImportAlias) or None + """ + self.aliases = aliases + self.names = [(a.name, a.asname) for a in aliases or []] + + def get_children(self): + yield from self.aliases or [] + class Attribute(NodeNG): """Class representing an :class:`ast.Attribute` node.""" @@ -3497,20 +3511,27 @@ def op_left_associative(self): return False -class Import(mixins.NoChildrenMixin, mixins.ImportFromMixin, Statement): - """Class representing an :class:`ast.Import` node. +class ImportAlias(mixins.NoChildrenMixin, NodeNG): + """Class representing an :class:`ast.alias` node. >>> node = astroid.extract_node('import astroid') >>> node + >>> node.aliases[0] + """ - _other_fields = ("names",) + _other_fields = ("name", "asname") - def __init__(self, names=None, lineno=None, col_offset=None, parent=None): + def __init__( + self, name=None, asname=None, lineno=None, col_offset=None, parent=None + ): """ - :param names: The names being imported. - :type names: list(tuple(str, str or None)) or None + :param name: The name being imported. + :type name: str or None + + :param asname: The alias of the imported name. + :type asname: str or None :param lineno: The line that this node appears on in the source code. :type lineno: int or None @@ -3522,18 +3543,56 @@ def __init__(self, names=None, lineno=None, col_offset=None, parent=None): :param parent: The parent node in the syntax tree. :type parent: NodeNG or None """ - self.names = names + self.name = name """The names being imported. - Each entry is a :class:`tuple` of the name being imported, - and the alias that the name is assigned to (if any). + :type: str or None + """ + self.asname = asname + """The alias of the imported name. - :type: list(tuple(str, str or None)) or None + :type: str or None """ super().__init__(lineno, col_offset, parent) +class Import(mixins.ImportFromMixin, Statement): + """Class representing an :class:`ast.Import` node. + + >>> node = astroid.extract_node('import astroid') + >>> node + + """ + + _astroid_fields = ("aliases",) + aliases = None + """ + :type: list(ImportAlias) or None + """ + _other_fields = ("names",) + names = None + """The names being imported. + + Each entry is a :class:`tuple` of the name being imported, + and the alias that the name is assigned to (if any). + + :type: list(tuple(str, str or None)) or None + """ + + def postinit(self, aliases=None): + """Do some setup after initialisation. + + :param aliases: The names being imported. + :type aliases: list(ImportAlias) or None + """ + self.aliases = aliases + self.names = [(a.name, a.asname) for a in aliases or []] + + def get_children(self): + yield from self.aliases or [] + + class Index(NodeNG): """Class representing an :class:`ast.Index` node. diff --git a/astroid/raw_building.py b/astroid/raw_building.py index 24fdc06950..bae60edbb1 100644 --- a/astroid/raw_building.py +++ b/astroid/raw_building.py @@ -100,7 +100,8 @@ def attach_import_node(node, modname, membername): """create a ImportFrom node and register it in the locals of the given node with the specified name """ - from_node = nodes.ImportFrom(modname, [(membername, None)]) + from_node = nodes.ImportFrom(modname) + from_node.postinit([nodes.ImportAlias(membername, None)]) _attach_local_node(node, from_node, membername) @@ -153,7 +154,9 @@ def build_function(name, args=None, posonlyargs=None, defaults=None, doc=None): def build_from_import(fromname, names): """create and initialize an astroid ImportFrom import statement""" - return nodes.ImportFrom(fromname, [(name, None) for name in names]) + node = nodes.ImportFrom(fromname) + node.postinit([nodes.ImportAlias(name, None) for name in names]) + return node def register_arguments(func, args=None): diff --git a/astroid/rebuilder.py b/astroid/rebuilder.py index 00d0d8c58d..3e152acd4e 100644 --- a/astroid/rebuilder.py +++ b/astroid/rebuilder.py @@ -955,15 +955,25 @@ def visit_importfrom( self, node: "ast.ImportFrom", parent: NodeNG ) -> nodes.ImportFrom: """visit an ImportFrom node by returning a fresh instance of it""" - names = [(alias.name, alias.asname) for alias in node.names] newnode = nodes.ImportFrom( node.module or "", - names, node.level or None, node.lineno, node.col_offset, parent, ) + newnode.postinit( + [ + nodes.ImportAlias( + alias.name, + alias.asname, + getattr(alias, "lineno", None), + getattr(alias, "col_offset", None), + newnode, + ) + for alias in node.names + ] + ) # store From names to add them to locals after building self._import_from_nodes.append(newnode) return newnode @@ -1100,13 +1110,24 @@ def visit_ifexp(self, node: "ast.IfExp", parent: NodeNG) -> nodes.IfExp: def visit_import(self, node: "ast.Import", parent: NodeNG) -> nodes.Import: """visit a Import node by returning a fresh instance of it""" - names = [(alias.name, alias.asname) for alias in node.names] newnode = nodes.Import( names, node.lineno, node.col_offset, parent, ) + newnode.postinit( + [ + nodes.ImportAlias( + alias.name, + alias.asname, + getattr(alias, "lineno", None), + getattr(alias, "col_offset", None), + newnode, + ) + for alias in node.names + ] + ) # save import names in parent's locals: for (name, asname) in newnode.names: name = asname or name