def import_to_node_single(imp: SortableImport, module: cst.Module) -> cst.BaseStatement: leading_lines = [ cst.EmptyLine(indent=True, comment=cst.Comment(line)) if line.startswith("#") else cst.EmptyLine(indent=False) for line in imp.comments.before ] trailing_whitespace = cst.TrailingWhitespace() trailing_comments = list(imp.comments.first_inline) names: List[cst.ImportAlias] = [] for item in imp.items: name = name_to_node(item.name) asname = cst.AsName( name=cst.Name(item.asname)) if item.asname else None node = cst.ImportAlias(name=name, asname=asname) names.append(node) trailing_comments += item.comments.before trailing_comments += item.comments.inline trailing_comments += item.comments.following trailing_comments += imp.comments.final trailing_comments += imp.comments.last_inline if trailing_comments: text = COMMENT_INDENT.join(trailing_comments) trailing_whitespace = cst.TrailingWhitespace( whitespace=cst.SimpleWhitespace(COMMENT_INDENT), comment=cst.Comment(text)) if imp.stem: stem, ndots = split_relative(imp.stem) if not stem: module_name = None else: module_name = name_to_node(stem) relative = (cst.Dot(), ) * ndots line = cst.SimpleStatementLine( body=[ cst.ImportFrom(module=module_name, names=names, relative=relative) ], leading_lines=leading_lines, trailing_whitespace=trailing_whitespace, ) else: line = cst.SimpleStatementLine( body=[cst.Import(names=names)], leading_lines=leading_lines, trailing_whitespace=trailing_whitespace, ) return line
def import_to_node_multi(imp: SortableImport, module: cst.Module) -> cst.BaseStatement: body: List[cst.BaseSmallStatement] = [] names: List[cst.ImportAlias] = [] prev: Optional[cst.ImportAlias] = None following: List[str] = [] lpar_lines: List[cst.EmptyLine] = [] lpar_inline: cst.TrailingWhitespace = cst.TrailingWhitespace() item_count = len(imp.items) for idx, item in enumerate(imp.items): name = name_to_node(item.name) asname = cst.AsName( name=cst.Name(item.asname)) if item.asname else None # Leading comments actually have to be trailing comments on the previous node. # That means putting them on the lpar node for the first item if item.comments.before: lines = [ cst.EmptyLine( indent=True, comment=cst.Comment(c), whitespace=cst.SimpleWhitespace(module.default_indent), ) for c in item.comments.before ] if prev is None: lpar_lines.extend(lines) else: prev.comma.whitespace_after.empty_lines.extend( lines) # type: ignore # all items except the last needs whitespace to indent the *next* line/item indent = idx != (len(imp.items) - 1) first_line = cst.TrailingWhitespace() inline = COMMENT_INDENT.join(item.comments.inline) if inline: first_line = cst.TrailingWhitespace( whitespace=cst.SimpleWhitespace(COMMENT_INDENT), comment=cst.Comment(inline), ) if idx == item_count - 1: following = item.comments.following + imp.comments.final else: following = item.comments.following after = cst.ParenthesizedWhitespace( indent=True, first_line=first_line, empty_lines=[ cst.EmptyLine( indent=True, comment=cst.Comment(c), whitespace=cst.SimpleWhitespace(module.default_indent), ) for c in following ], last_line=cst.SimpleWhitespace( module.default_indent if indent else ""), ) node = cst.ImportAlias( name=name, asname=asname, comma=cst.Comma(whitespace_after=after), ) names.append(node) prev = node # from foo import ( # bar # ) if imp.stem: stem, ndots = split_relative(imp.stem) if not stem: module_name = None else: module_name = name_to_node(stem) relative = (cst.Dot(), ) * ndots # inline comment following lparen if imp.comments.first_inline: inline = COMMENT_INDENT.join(imp.comments.first_inline) lpar_inline = cst.TrailingWhitespace( whitespace=cst.SimpleWhitespace(COMMENT_INDENT), comment=cst.Comment(inline), ) body = [ cst.ImportFrom( module=module_name, names=names, relative=relative, lpar=cst.LeftParen( whitespace_after=cst.ParenthesizedWhitespace( indent=True, first_line=lpar_inline, empty_lines=lpar_lines, last_line=cst.SimpleWhitespace(module.default_indent), ), ), rpar=cst.RightParen(), ) ] # import foo else: raise ValueError("can't render basic imports on multiple lines") # comment lines above import leading_lines = [ cst.EmptyLine(indent=True, comment=cst.Comment(line)) if line.startswith("#") else cst.EmptyLine(indent=False) for line in imp.comments.before ] # inline comments following import/rparen if imp.comments.last_inline: inline = COMMENT_INDENT.join(imp.comments.last_inline) trailing = cst.TrailingWhitespace( whitespace=cst.SimpleWhitespace(COMMENT_INDENT), comment=cst.Comment(inline)) else: trailing = cst.TrailingWhitespace() return cst.SimpleStatementLine( body=body, leading_lines=leading_lines, trailing_whitespace=trailing, )
class ImportFromParseTest(CSTNodeTest): @data_provider( ( # Simple from import statement { "node": cst.ImportFrom( module=cst.Name("foo"), names=(cst.ImportAlias(cst.Name("bar")),) ), "code": "from foo import bar", }, # From import statement with alias { "node": cst.ImportFrom( module=cst.Name("foo"), names=( cst.ImportAlias( cst.Name("bar"), asname=cst.AsName(cst.Name("baz")) ), ), ), "code": "from foo import bar as baz", }, # Multiple imports { "node": cst.ImportFrom( module=cst.Name("foo"), names=( cst.ImportAlias( cst.Name("bar"), comma=cst.Comma(whitespace_after=cst.SimpleWhitespace(" ")), ), cst.ImportAlias(cst.Name("baz")), ), ), "code": "from foo import bar, baz", }, # Trailing comma { "node": cst.ImportFrom( module=cst.Name("foo"), names=( cst.ImportAlias( cst.Name("bar"), comma=cst.Comma(whitespace_after=cst.SimpleWhitespace(" ")), ), cst.ImportAlias(cst.Name("baz"), comma=cst.Comma()), ), ), "code": "from foo import bar, baz,", }, # Star import statement { "node": cst.ImportFrom(module=cst.Name("foo"), names=cst.ImportStar()), "code": "from foo import *", }, # Simple relative import statement { "node": cst.ImportFrom( relative=(cst.Dot(),), module=cst.Name("foo"), names=(cst.ImportAlias(cst.Name("bar")),), ), "code": "from .foo import bar", }, { "node": cst.ImportFrom( relative=(cst.Dot(), cst.Dot()), module=cst.Name("foo"), names=(cst.ImportAlias(cst.Name("bar")),), ), "code": "from ..foo import bar", }, # Relative only import { "node": cst.ImportFrom( relative=(cst.Dot(), cst.Dot()), module=None, names=(cst.ImportAlias(cst.Name("bar")),), ), "code": "from .. import bar", }, # Parenthesis { "node": cst.ImportFrom( module=cst.Name("foo"), lpar=cst.LeftParen(), names=( cst.ImportAlias( cst.Name("bar"), asname=cst.AsName(cst.Name("baz")) ), ), rpar=cst.RightParen(), ), "code": "from foo import (bar as baz)", }, # Verify whitespace works everywhere. { "node": cst.ImportFrom( relative=( cst.Dot( whitespace_before=cst.SimpleWhitespace(""), whitespace_after=cst.SimpleWhitespace(" "), ), cst.Dot( whitespace_before=cst.SimpleWhitespace(""), whitespace_after=cst.SimpleWhitespace(" "), ), ), module=cst.Name("foo"), lpar=cst.LeftParen(whitespace_after=cst.SimpleWhitespace(" ")), names=( cst.ImportAlias( cst.Name("bar"), asname=cst.AsName( cst.Name("baz"), whitespace_before_as=cst.SimpleWhitespace(" "), whitespace_after_as=cst.SimpleWhitespace(" "), ), comma=cst.Comma( whitespace_before=cst.SimpleWhitespace(" "), whitespace_after=cst.SimpleWhitespace(" "), ), ), cst.ImportAlias( cst.Name("unittest"), asname=cst.AsName( cst.Name("ut"), whitespace_before_as=cst.SimpleWhitespace(" "), whitespace_after_as=cst.SimpleWhitespace(" "), ), ), ), rpar=cst.RightParen(whitespace_before=cst.SimpleWhitespace(" ")), whitespace_after_from=cst.SimpleWhitespace(" "), whitespace_before_import=cst.SimpleWhitespace(" "), whitespace_after_import=cst.SimpleWhitespace(" "), ), "code": "from . . foo import ( bar as baz , unittest as ut )", }, ) ) def test_valid(self, **kwargs: Any) -> None: self.validate_node( parser=lambda code: ensure_type( parse_statement(code), cst.SimpleStatementLine ).body[0], **kwargs, )
class ImportFromCreateTest(CSTNodeTest): @data_provider( ( # Simple from import statement { "node": cst.ImportFrom( module=cst.Name("foo"), names=(cst.ImportAlias(cst.Name("bar")),) ), "code": "from foo import bar", }, # From import statement with alias { "node": cst.ImportFrom( module=cst.Name("foo"), names=( cst.ImportAlias( cst.Name("bar"), asname=cst.AsName(cst.Name("baz")) ), ), ), "code": "from foo import bar as baz", }, # Multiple imports { "node": cst.ImportFrom( module=cst.Name("foo"), names=( cst.ImportAlias(cst.Name("bar")), cst.ImportAlias(cst.Name("baz")), ), ), "code": "from foo import bar, baz", }, # Trailing comma { "node": cst.ImportFrom( module=cst.Name("foo"), names=( cst.ImportAlias(cst.Name("bar"), comma=cst.Comma()), cst.ImportAlias(cst.Name("baz"), comma=cst.Comma()), ), ), "code": "from foo import bar,baz,", "expected_position": CodeRange((1, 0), (1, 23)), }, # Star import statement { "node": cst.ImportFrom(module=cst.Name("foo"), names=cst.ImportStar()), "code": "from foo import *", "expected_position": CodeRange((1, 0), (1, 17)), }, # Simple relative import statement { "node": cst.ImportFrom( relative=(cst.Dot(),), module=cst.Name("foo"), names=(cst.ImportAlias(cst.Name("bar")),), ), "code": "from .foo import bar", }, { "node": cst.ImportFrom( relative=(cst.Dot(), cst.Dot()), module=cst.Name("foo"), names=(cst.ImportAlias(cst.Name("bar")),), ), "code": "from ..foo import bar", }, # Relative only import { "node": cst.ImportFrom( relative=(cst.Dot(), cst.Dot()), module=None, names=(cst.ImportAlias(cst.Name("bar")),), ), "code": "from .. import bar", }, # Parenthesis { "node": cst.ImportFrom( module=cst.Name("foo"), lpar=cst.LeftParen(), names=( cst.ImportAlias( cst.Name("bar"), asname=cst.AsName(cst.Name("baz")) ), ), rpar=cst.RightParen(), ), "code": "from foo import (bar as baz)", "expected_position": CodeRange((1, 0), (1, 28)), }, # Verify whitespace works everywhere. { "node": cst.ImportFrom( relative=( cst.Dot( whitespace_before=cst.SimpleWhitespace(" "), whitespace_after=cst.SimpleWhitespace(" "), ), cst.Dot( whitespace_before=cst.SimpleWhitespace(" "), whitespace_after=cst.SimpleWhitespace(" "), ), ), module=cst.Name("foo"), lpar=cst.LeftParen(whitespace_after=cst.SimpleWhitespace(" ")), names=( cst.ImportAlias( cst.Name("bar"), asname=cst.AsName( cst.Name("baz"), whitespace_before_as=cst.SimpleWhitespace(" "), whitespace_after_as=cst.SimpleWhitespace(" "), ), comma=cst.Comma( whitespace_before=cst.SimpleWhitespace(" "), whitespace_after=cst.SimpleWhitespace(" "), ), ), cst.ImportAlias( cst.Name("unittest"), asname=cst.AsName( cst.Name("ut"), whitespace_before_as=cst.SimpleWhitespace(" "), whitespace_after_as=cst.SimpleWhitespace(" "), ), ), ), rpar=cst.RightParen(whitespace_before=cst.SimpleWhitespace(" ")), whitespace_after_from=cst.SimpleWhitespace(" "), whitespace_before_import=cst.SimpleWhitespace(" "), whitespace_after_import=cst.SimpleWhitespace(" "), ), "code": "from . . foo import ( bar as baz , unittest as ut )", "expected_position": CodeRange((1, 0), (1, 61)), }, ) ) def test_valid(self, **kwargs: Any) -> None: self.validate_node(**kwargs) @data_provider( ( { "get_node": lambda: cst.ImportFrom( module=None, names=(cst.ImportAlias(cst.Name("bar")),) ), "expected_re": "Must have a module specified", }, { "get_node": lambda: cst.ImportFrom(module=cst.Name("foo"), names=()), "expected_re": "at least one ImportAlias", }, { "get_node": lambda: cst.ImportFrom( module=cst.Name("foo"), names=(cst.ImportAlias(cst.Name("bar")),), lpar=cst.LeftParen(), ), "expected_re": "left paren without right paren", }, { "get_node": lambda: cst.ImportFrom( module=cst.Name("foo"), names=(cst.ImportAlias(cst.Name("bar")),), rpar=cst.RightParen(), ), "expected_re": "right paren without left paren", }, { "get_node": lambda: cst.ImportFrom( module=cst.Name("foo"), names=cst.ImportStar(), lpar=cst.LeftParen() ), "expected_re": "cannot have parens", }, { "get_node": lambda: cst.ImportFrom( module=cst.Name("foo"), names=cst.ImportStar(), rpar=cst.RightParen(), ), "expected_re": "cannot have parens", }, { "get_node": lambda: cst.ImportFrom( module=cst.Name("foo"), names=(cst.ImportAlias(cst.Name("bar")),), whitespace_after_from=cst.SimpleWhitespace(""), ), "expected_re": "one space after from", }, { "get_node": lambda: cst.ImportFrom( module=cst.Name("foo"), names=(cst.ImportAlias(cst.Name("bar")),), whitespace_before_import=cst.SimpleWhitespace(""), ), "expected_re": "one space before import", }, { "get_node": lambda: cst.ImportFrom( module=cst.Name("foo"), names=(cst.ImportAlias(cst.Name("bar")),), whitespace_after_import=cst.SimpleWhitespace(""), ), "expected_re": "one space after import", }, ) ) def test_invalid(self, **kwargs: Any) -> None: self.assert_invalid(**kwargs)
class ImportParseTest(CSTNodeTest): @data_provider( ( # Simple import statement { "node": cst.Import(names=(cst.ImportAlias(cst.Name("foo")),)), "code": "import foo", }, { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")) ), ) ), "code": "import foo.bar", }, { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")) ), ) ), "code": "import foo.bar", }, # Comma-separated list of imports { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")), comma=cst.Comma(whitespace_after=cst.SimpleWhitespace(" ")), ), cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("baz")) ), ) ), "code": "import foo.bar, foo.baz", }, # Import with an alias { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")), asname=cst.AsName(cst.Name("baz")), ), ) ), "code": "import foo.bar as baz", }, # Import with an alias, comma separated { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")), asname=cst.AsName(cst.Name("baz")), comma=cst.Comma(whitespace_after=cst.SimpleWhitespace(" ")), ), cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("baz")), asname=cst.AsName(cst.Name("bar")), ), ) ), "code": "import foo.bar as baz, foo.baz as bar", }, # Combine for fun and profit { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")), asname=cst.AsName(cst.Name("baz")), comma=cst.Comma(whitespace_after=cst.SimpleWhitespace(" ")), ), cst.ImportAlias( cst.Attribute(cst.Name("insta"), cst.Name("gram")), comma=cst.Comma(whitespace_after=cst.SimpleWhitespace(" ")), ), cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("baz")), comma=cst.Comma(whitespace_after=cst.SimpleWhitespace(" ")), ), cst.ImportAlias( cst.Name("unittest"), asname=cst.AsName(cst.Name("ut")) ), ) ), "code": "import foo.bar as baz, insta.gram, foo.baz, unittest as ut", }, # Verify whitespace works everywhere. { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute( cst.Name("foo"), cst.Name("bar"), dot=cst.Dot( whitespace_before=cst.SimpleWhitespace(" "), whitespace_after=cst.SimpleWhitespace(" "), ), ), asname=cst.AsName( cst.Name("baz"), whitespace_before_as=cst.SimpleWhitespace(" "), whitespace_after_as=cst.SimpleWhitespace(" "), ), comma=cst.Comma( whitespace_before=cst.SimpleWhitespace(" "), whitespace_after=cst.SimpleWhitespace(" "), ), ), cst.ImportAlias( cst.Name("unittest"), asname=cst.AsName( cst.Name("ut"), whitespace_before_as=cst.SimpleWhitespace(" "), whitespace_after_as=cst.SimpleWhitespace(" "), ), ), ), whitespace_after_import=cst.SimpleWhitespace(" "), ), "code": "import foo . bar as baz , unittest as ut", }, ) ) def test_valid(self, **kwargs: Any) -> None: self.validate_node( parser=lambda code: ensure_type( parse_statement(code), cst.SimpleStatementLine ).body[0], **kwargs, )
class ImportCreateTest(CSTNodeTest): @data_provider( ( # Simple import statement { "node": cst.Import(names=(cst.ImportAlias(cst.Name("foo")),)), "code": "import foo", }, { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")) ), ) ), "code": "import foo.bar", }, { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")) ), ) ), "code": "import foo.bar", }, # Comma-separated list of imports { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")) ), cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("baz")) ), ) ), "code": "import foo.bar, foo.baz", "expected_position": CodeRange((1, 0), (1, 23)), }, # Import with an alias { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")), asname=cst.AsName(cst.Name("baz")), ), ) ), "code": "import foo.bar as baz", }, # Import with an alias, comma separated { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")), asname=cst.AsName(cst.Name("baz")), ), cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("baz")), asname=cst.AsName(cst.Name("bar")), ), ) ), "code": "import foo.bar as baz, foo.baz as bar", }, # Combine for fun and profit { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")), asname=cst.AsName(cst.Name("baz")), ), cst.ImportAlias( cst.Attribute(cst.Name("insta"), cst.Name("gram")) ), cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("baz")) ), cst.ImportAlias( cst.Name("unittest"), asname=cst.AsName(cst.Name("ut")) ), ) ), "code": "import foo.bar as baz, insta.gram, foo.baz, unittest as ut", }, # Verify whitespace works everywhere. { "node": cst.Import( names=( cst.ImportAlias( cst.Attribute( cst.Name("foo"), cst.Name("bar"), dot=cst.Dot( whitespace_before=cst.SimpleWhitespace(" "), whitespace_after=cst.SimpleWhitespace(" "), ), ), asname=cst.AsName( cst.Name("baz"), whitespace_before_as=cst.SimpleWhitespace(" "), whitespace_after_as=cst.SimpleWhitespace(" "), ), comma=cst.Comma( whitespace_before=cst.SimpleWhitespace(" "), whitespace_after=cst.SimpleWhitespace(" "), ), ), cst.ImportAlias( cst.Name("unittest"), asname=cst.AsName( cst.Name("ut"), whitespace_before_as=cst.SimpleWhitespace(" "), whitespace_after_as=cst.SimpleWhitespace(" "), ), ), ), whitespace_after_import=cst.SimpleWhitespace(" "), ), "code": "import foo . bar as baz , unittest as ut", "expected_position": CodeRange((1, 0), (1, 46)), }, ) ) def test_valid(self, **kwargs: Any) -> None: self.validate_node(**kwargs) @data_provider( ( { "get_node": lambda: cst.Import(names=()), "expected_re": "at least one ImportAlias", }, { "get_node": lambda: cst.Import(names=(cst.ImportAlias(cst.Name("")),)), "expected_re": "empty name identifier", }, { "get_node": lambda: cst.Import( names=( cst.ImportAlias(cst.Attribute(cst.Name(""), cst.Name("bla"))), ) ), "expected_re": "empty name identifier", }, { "get_node": lambda: cst.Import( names=( cst.ImportAlias(cst.Attribute(cst.Name("bla"), cst.Name(""))), ) ), "expected_re": "empty name identifier", }, { "get_node": lambda: cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")), comma=cst.Comma(), ), ) ), "expected_re": "trailing comma", }, { "get_node": lambda: cst.Import( names=( cst.ImportAlias( cst.Attribute(cst.Name("foo"), cst.Name("bar")) ), ), whitespace_after_import=cst.SimpleWhitespace(""), ), "expected_re": "at least one space", }, ) ) def test_invalid(self, **kwargs: Any) -> None: self.assert_invalid(**kwargs)
class AttributeTest(CSTNodeTest): @data_provider( ( # Simple attribute access { "node": cst.Attribute(cst.Name("foo"), cst.Name("bar")), "code": "foo.bar", "parser": parse_expression, "expected_position": CodeRange((1, 0), (1, 7)), }, # Parenthesized attribute access { "node": cst.Attribute( lpar=(cst.LeftParen(),), value=cst.Name("foo"), attr=cst.Name("bar"), rpar=(cst.RightParen(),), ), "code": "(foo.bar)", "parser": parse_expression, "expected_position": CodeRange((1, 1), (1, 8)), }, # Make sure that spacing works { "node": cst.Attribute( lpar=(cst.LeftParen(whitespace_after=cst.SimpleWhitespace(" ")),), value=cst.Name("foo"), dot=cst.Dot( whitespace_before=cst.SimpleWhitespace(" "), whitespace_after=cst.SimpleWhitespace(" "), ), attr=cst.Name("bar"), rpar=(cst.RightParen(whitespace_before=cst.SimpleWhitespace(" ")),), ), "code": "( foo . bar )", "parser": parse_expression, "expected_position": CodeRange((1, 2), (1, 11)), }, ) ) def test_valid(self, **kwargs: Any) -> None: self.validate_node(**kwargs) @data_provider( ( { "get_node": ( lambda: cst.Attribute( cst.Name("foo"), cst.Name("bar"), lpar=(cst.LeftParen(),) ) ), "expected_re": "left paren without right paren", }, { "get_node": ( lambda: cst.Attribute( cst.Name("foo"), cst.Name("bar"), rpar=(cst.RightParen(),) ) ), "expected_re": "right paren without left paren", }, ) ) def test_invalid(self, **kwargs: Any) -> None: self.assert_invalid(**kwargs)