def leave_ImportFrom(self, original_node: libcst.ImportFrom, updated_node: libcst.ImportFrom) -> libcst.ImportFrom: if isinstance(updated_node.names, libcst.ImportStar): # There's nothing to do here! return updated_node # Get the module we're importing as a string, see if we have work to do. module = get_absolute_module_for_import(self.context.full_module_name, updated_node) if (module is None or module not in self.module_mapping and module not in self.alias_mapping): return updated_node # We have work to do, mark that we won't modify this again. imports_to_add = self.module_mapping.get(module, []) if module in self.module_mapping: del self.module_mapping[module] aliases_to_add = self.alias_mapping.get(module, []) if module in self.alias_mapping: del self.alias_mapping[module] # Now, do the actual update. return updated_node.with_changes(names=[ *(libcst.ImportAlias(name=libcst.Name(imp)) for imp in sorted(imports_to_add)), *(libcst.ImportAlias( name=libcst.Name(imp), asname=libcst.AsName(name=libcst.Name(alias)), ) for (imp, alias) in sorted(aliases_to_add)), *updated_node.names, ])
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 leave_ImportFrom(self, original_node: libcst.ImportFrom, updated_node: libcst.ImportFrom) -> libcst.ImportFrom: if len(updated_node.relative) > 0 or updated_node.module is None: # Don't support relative-only imports at the moment. return updated_node if updated_node.names == "*": # There's nothing to do here! return updated_node # Get the module we're importing as a string, see if we have work to do module = self._get_string_name(updated_node.module) if module not in self.module_mapping and module not in self.alias_mapping: return updated_node # We have work to do, mark that we won't modify this again. imports_to_add = self.module_mapping.get(module, []) if module in self.module_mapping: del self.module_mapping[module] aliases_to_add = self.alias_mapping.get(module, []) if module in self.alias_mapping: del self.alias_mapping[module] # Now, do the actual update. return updated_node.with_changes(names=( *[ libcst.ImportAlias(name=libcst.Name(imp)) for imp in imports_to_add ], *[ libcst.ImportAlias( name=libcst.Name(imp), asname=libcst.AsName(name=libcst.Name(alias)), ) for (imp, alias) in aliases_to_add ], *updated_node.names, ))
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 TryTest(CSTNodeTest): @data_provider( ( # Simple try/except block { "node": cst.Try( cst.SimpleStatementSuite((cst.Pass(),)), handlers=( cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), whitespace_after_except=cst.SimpleWhitespace(""), ), ), ), "code": "try: pass\nexcept: pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (2, 12)), }, # Try/except with a class { "node": cst.Try( cst.SimpleStatementSuite((cst.Pass(),)), handlers=( cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), type=cst.Name("Exception"), ), ), ), "code": "try: pass\nexcept Exception: pass\n", "parser": parse_statement, }, # Try/except with a named class { "node": cst.Try( cst.SimpleStatementSuite((cst.Pass(),)), handlers=( cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), type=cst.Name("Exception"), name=cst.AsName(cst.Name("exc")), ), ), ), "code": "try: pass\nexcept Exception as exc: pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (2, 29)), }, # Try/except with multiple clauses { "node": cst.Try( cst.SimpleStatementSuite((cst.Pass(),)), handlers=( cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), type=cst.Name("TypeError"), name=cst.AsName(cst.Name("e")), ), cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), type=cst.Name("KeyError"), name=cst.AsName(cst.Name("e")), ), cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), whitespace_after_except=cst.SimpleWhitespace(""), ), ), ), "code": "try: pass\n" + "except TypeError as e: pass\n" + "except KeyError as e: pass\n" + "except: pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (4, 12)), }, # Simple try/finally block { "node": cst.Try( cst.SimpleStatementSuite((cst.Pass(),)), finalbody=cst.Finally(cst.SimpleStatementSuite((cst.Pass(),))), ), "code": "try: pass\nfinally: pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (2, 13)), }, # Simple try/except/finally block { "node": cst.Try( cst.SimpleStatementSuite((cst.Pass(),)), handlers=( cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), whitespace_after_except=cst.SimpleWhitespace(""), ), ), finalbody=cst.Finally(cst.SimpleStatementSuite((cst.Pass(),))), ), "code": "try: pass\nexcept: pass\nfinally: pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (3, 13)), }, # Simple try/except/else block { "node": cst.Try( cst.SimpleStatementSuite((cst.Pass(),)), handlers=( cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), whitespace_after_except=cst.SimpleWhitespace(""), ), ), orelse=cst.Else(cst.SimpleStatementSuite((cst.Pass(),))), ), "code": "try: pass\nexcept: pass\nelse: pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (3, 10)), }, # Simple try/except/else block/finally { "node": cst.Try( cst.SimpleStatementSuite((cst.Pass(),)), handlers=( cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), whitespace_after_except=cst.SimpleWhitespace(""), ), ), orelse=cst.Else(cst.SimpleStatementSuite((cst.Pass(),))), finalbody=cst.Finally(cst.SimpleStatementSuite((cst.Pass(),))), ), "code": "try: pass\nexcept: pass\nelse: pass\nfinally: pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (4, 13)), }, # Verify whitespace in various locations { "node": cst.Try( leading_lines=(cst.EmptyLine(comment=cst.Comment("# 1")),), body=cst.SimpleStatementSuite((cst.Pass(),)), handlers=( cst.ExceptHandler( leading_lines=(cst.EmptyLine(comment=cst.Comment("# 2")),), type=cst.Name("TypeError"), name=cst.AsName( cst.Name("e"), whitespace_before_as=cst.SimpleWhitespace(" "), whitespace_after_as=cst.SimpleWhitespace(" "), ), whitespace_after_except=cst.SimpleWhitespace(" "), whitespace_before_colon=cst.SimpleWhitespace(" "), body=cst.SimpleStatementSuite((cst.Pass(),)), ), ), orelse=cst.Else( leading_lines=(cst.EmptyLine(comment=cst.Comment("# 3")),), body=cst.SimpleStatementSuite((cst.Pass(),)), whitespace_before_colon=cst.SimpleWhitespace(" "), ), finalbody=cst.Finally( leading_lines=(cst.EmptyLine(comment=cst.Comment("# 4")),), body=cst.SimpleStatementSuite((cst.Pass(),)), whitespace_before_colon=cst.SimpleWhitespace(" "), ), whitespace_before_colon=cst.SimpleWhitespace(" "), ), "code": "# 1\ntry : pass\n# 2\nexcept TypeError as e : pass\n# 3\nelse : pass\n# 4\nfinally : pass\n", "parser": parse_statement, "expected_position": CodeRange((2, 0), (8, 14)), }, # Please don't write code like this { "node": cst.Try( cst.SimpleStatementSuite((cst.Pass(),)), handlers=( cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), type=cst.Name("TypeError"), name=cst.AsName(cst.Name("e")), ), cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), type=cst.Name("KeyError"), name=cst.AsName(cst.Name("e")), ), cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), whitespace_after_except=cst.SimpleWhitespace(""), ), ), orelse=cst.Else(cst.SimpleStatementSuite((cst.Pass(),))), finalbody=cst.Finally(cst.SimpleStatementSuite((cst.Pass(),))), ), "code": "try: pass\n" + "except TypeError as e: pass\n" + "except KeyError as e: pass\n" + "except: pass\n" + "else: pass\n" + "finally: pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (6, 13)), }, # Verify indentation { "node": DummyIndentedBlock( " ", cst.Try( cst.SimpleStatementSuite((cst.Pass(),)), handlers=( cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), type=cst.Name("TypeError"), name=cst.AsName(cst.Name("e")), ), cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), type=cst.Name("KeyError"), name=cst.AsName(cst.Name("e")), ), cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), whitespace_after_except=cst.SimpleWhitespace(""), ), ), orelse=cst.Else(cst.SimpleStatementSuite((cst.Pass(),))), finalbody=cst.Finally(cst.SimpleStatementSuite((cst.Pass(),))), ), ), "code": " try: pass\n" + " except TypeError as e: pass\n" + " except KeyError as e: pass\n" + " except: pass\n" + " else: pass\n" + " finally: pass\n", "parser": None, }, # Verify indentation in bodies { "node": DummyIndentedBlock( " ", cst.Try( cst.IndentedBlock((cst.SimpleStatementLine((cst.Pass(),)),)), handlers=( cst.ExceptHandler( cst.IndentedBlock( (cst.SimpleStatementLine((cst.Pass(),)),) ), whitespace_after_except=cst.SimpleWhitespace(""), ), ), orelse=cst.Else( cst.IndentedBlock((cst.SimpleStatementLine((cst.Pass(),)),)) ), finalbody=cst.Finally( cst.IndentedBlock((cst.SimpleStatementLine((cst.Pass(),)),)) ), ), ), "code": " try:\n" + " pass\n" + " except:\n" + " pass\n" + " else:\n" + " pass\n" + " finally:\n" + " pass\n", "parser": None, }, # No space when using grouping parens { "node": cst.Try( cst.SimpleStatementSuite((cst.Pass(),)), handlers=( cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), whitespace_after_except=cst.SimpleWhitespace(""), type=cst.Name( "Exception", lpar=(cst.LeftParen(),), rpar=(cst.RightParen(),), ), ), ), ), "code": "try: pass\nexcept(Exception): pass\n", "parser": parse_statement, }, # No space when using tuple { "node": cst.Try( cst.SimpleStatementSuite((cst.Pass(),)), handlers=( cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), whitespace_after_except=cst.SimpleWhitespace(""), type=cst.Tuple( [ cst.Element( cst.Name("IOError"), comma=cst.Comma( whitespace_after=cst.SimpleWhitespace(" ") ), ), cst.Element(cst.Name("ImportError")), ] ), ), ), ), "code": "try: pass\nexcept(IOError, ImportError): pass\n", "parser": parse_statement, }, ) ) def test_valid(self, **kwargs: Any) -> None: self.validate_node(**kwargs) @data_provider( ( { "get_node": lambda: cst.AsName(cst.Name("")), "expected_re": "empty name identifier", }, { "get_node": lambda: cst.AsName( cst.Name("bla"), whitespace_after_as=cst.SimpleWhitespace("") ), "expected_re": "between 'as'", }, { "get_node": lambda: cst.AsName( cst.Name("bla"), whitespace_before_as=cst.SimpleWhitespace("") ), "expected_re": "before 'as'", }, { "get_node": lambda: cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), name=cst.AsName(cst.Name("bla")), ), "expected_re": "name for an empty type", }, { "get_node": lambda: cst.ExceptHandler( cst.SimpleStatementSuite((cst.Pass(),)), type=cst.Name("TypeError"), whitespace_after_except=cst.SimpleWhitespace(""), ), "expected_re": "at least one space after except", }, { "get_node": lambda: cst.Try(cst.SimpleStatementSuite((cst.Pass(),))), "expected_re": "at least one ExceptHandler or Finally", }, { "get_node": lambda: cst.Try( cst.SimpleStatementSuite((cst.Pass(),)), orelse=cst.Else(cst.SimpleStatementSuite((cst.Pass(),))), finalbody=cst.Finally(cst.SimpleStatementSuite((cst.Pass(),))), ), "expected_re": "at least one ExceptHandler in order to have an Else", }, ) ) def test_invalid(self, **kwargs: Any) -> None: self.assert_invalid(**kwargs)
class WithTest(CSTNodeTest): @data_provider(( # Simple with block { "node": cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), ), "code": "with context_mgr(): pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (1, 24)), }, # Simple async with block { "node": cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), asynchronous=cst.Asynchronous(), ), "code": "async with context_mgr(): pass\n", "parser": lambda code: parse_statement( code, config=PartialParserConfig(python_version="3.7")), }, # Python 3.6 async with block { "node": cst.FunctionDef( cst.Name("foo"), cst.Parameters(), cst.IndentedBlock((cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), asynchronous=cst.Asynchronous(), ), )), asynchronous=cst.Asynchronous(), ), "code": "async def foo():\n async with context_mgr(): pass\n", "parser": lambda code: parse_statement( code, config=PartialParserConfig(python_version="3.6")), }, # Multiple context managers { "node": cst.With( ( cst.WithItem(cst.Call(cst.Name("foo"))), cst.WithItem(cst.Call(cst.Name("bar"))), ), cst.SimpleStatementSuite((cst.Pass(), )), ), "code": "with foo(), bar(): pass\n", "parser": None, }, { "node": cst.With( ( cst.WithItem( cst.Call(cst.Name("foo")), comma=cst.Comma( whitespace_after=cst.SimpleWhitespace(" ")), ), cst.WithItem(cst.Call(cst.Name("bar"))), ), cst.SimpleStatementSuite((cst.Pass(), )), ), "code": "with foo(), bar(): pass\n", "parser": parse_statement, }, # With block containing variable for context manager. { "node": cst.With( (cst.WithItem( cst.Call(cst.Name("context_mgr")), cst.AsName(cst.Name("ctx")), ), ), cst.SimpleStatementSuite((cst.Pass(), )), ), "code": "with context_mgr() as ctx: pass\n", "parser": parse_statement, }, # indentation { "node": DummyIndentedBlock( " ", cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), ), ), "code": " with context_mgr(): pass\n", "parser": None, "expected_position": CodeRange((1, 4), (1, 28)), }, # with an indented body { "node": DummyIndentedBlock( " ", cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.IndentedBlock((cst.SimpleStatementLine( (cst.Pass(), )), )), ), ), "code": " with context_mgr():\n pass\n", "parser": None, "expected_position": CodeRange((1, 4), (2, 12)), }, # leading_lines { "node": cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), leading_lines=(cst.EmptyLine( comment=cst.Comment("# leading comment")), ), ), "code": "# leading comment\nwith context_mgr(): pass\n", "parser": parse_statement, "expected_position": CodeRange((2, 0), (2, 24)), }, # Weird spacing rules { "node": cst.With( (cst.WithItem( cst.Call( cst.Name("context_mgr"), lpar=(cst.LeftParen(), ), rpar=(cst.RightParen(), ), )), ), cst.SimpleStatementSuite((cst.Pass(), )), whitespace_after_with=cst.SimpleWhitespace(""), ), "code": "with(context_mgr()): pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (1, 25)), }, # Whitespace { "node": cst.With( (cst.WithItem( cst.Call(cst.Name("context_mgr")), cst.AsName( cst.Name("ctx"), whitespace_before_as=cst.SimpleWhitespace(" "), whitespace_after_as=cst.SimpleWhitespace(" "), ), ), ), cst.SimpleStatementSuite((cst.Pass(), )), whitespace_after_with=cst.SimpleWhitespace(" "), whitespace_before_colon=cst.SimpleWhitespace(" "), ), "code": "with context_mgr() as ctx : pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (1, 36)), }, )) def test_valid(self, **kwargs: Any) -> None: self.validate_node(**kwargs) @data_provider(( { "get_node": lambda: cst.With((), cst.IndentedBlock((cst.SimpleStatementLine( (cst.Pass(), )), ))), "expected_re": "A With statement must have at least one WithItem", }, { "get_node": lambda: cst.With( (cst.WithItem( cst.Call(cst.Name("foo")), comma=cst.Comma(whitespace_after=cst.SimpleWhitespace(" ") ), ), ), cst.IndentedBlock((cst.SimpleStatementLine((cst.Pass(), )), )), ), "expected_re": "The last WithItem in a With cannot have a trailing comma", }, { "get_node": lambda: cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), whitespace_after_with=cst.SimpleWhitespace(""), ), "expected_re": "Must have at least one space after with keyword", }, )) def test_invalid(self, **kwargs: Any) -> None: self.assert_invalid(**kwargs) @data_provider(( { "code": "with a, b: pass", "parser": parse_statement_as(python_version="3.1"), "expect_success": True, }, { "code": "with a, b: pass", "parser": parse_statement_as(python_version="3.0"), "expect_success": False, }, )) def test_versions(self, **kwargs: Any) -> None: self.assert_parses(**kwargs)
class StatementTest(UnitTest): @data_provider( ( # Simple imports that are already absolute. (None, "from a.b import c", "a.b"), ("x.y.z", "from a.b import c", "a.b"), # Relative import that can't be resolved due to missing module. (None, "from ..w import c", None), # Relative import that goes past the module level. ("x", "from ...y import z", None), ("x.y.z", "from .....w import c", None), ("x.y.z", "from ... import c", None), # Correct resolution of absolute from relative modules. ("x.y.z", "from . import c", "x.y"), ("x.y.z", "from .. import c", "x"), ("x.y.z", "from .w import c", "x.y.w"), ("x.y.z", "from ..w import c", "x.w"), ("x.y.z", "from ...w import c", "w"), ) ) def test_get_absolute_module( self, module: Optional[str], importfrom: str, output: Optional[str], ) -> None: node = ensure_type(cst.parse_statement(importfrom), cst.SimpleStatementLine) assert len(node.body) == 1, "Unexpected number of statements!" import_node = ensure_type(node.body[0], cst.ImportFrom) self.assertEqual(get_absolute_module_for_import(module, import_node), output) if output is None: with self.assertRaises(Exception): get_absolute_module_for_import_or_raise(module, import_node) else: self.assertEqual( get_absolute_module_for_import_or_raise(module, import_node), output ) @data_provider( ( # Nodes without an asname (cst.ImportAlias(name=cst.Name("foo")), "foo", None), ( cst.ImportAlias(name=cst.Attribute(cst.Name("foo"), cst.Name("bar"))), "foo.bar", None, ), # Nodes with an asname ( cst.ImportAlias( name=cst.Name("foo"), asname=cst.AsName(name=cst.Name("baz")) ), "foo", "baz", ), ( cst.ImportAlias( name=cst.Attribute(cst.Name("foo"), cst.Name("bar")), asname=cst.AsName(name=cst.Name("baz")), ), "foo.bar", "baz", ), ) ) def test_importalias_helpers( self, alias_node: cst.ImportAlias, full_name: str, alias: Optional[str] ) -> None: self.assertEqual(alias_node.evaluated_name, full_name) self.assertEqual(alias_node.evaluated_alias, alias)
class TryStarTest(CSTNodeTest): @data_provider(( # Try/except with a class { "node": cst.TryStar( cst.SimpleStatementSuite((cst.Pass(), )), handlers=(cst.ExceptStarHandler( cst.SimpleStatementSuite((cst.Pass(), )), type=cst.Name("Exception"), ), ), ), "code": "try: pass\nexcept* Exception: pass\n", "parser": native_parse_statement, }, # Try/except with a named class { "node": cst.TryStar( cst.SimpleStatementSuite((cst.Pass(), )), handlers=(cst.ExceptStarHandler( cst.SimpleStatementSuite((cst.Pass(), )), type=cst.Name("Exception"), name=cst.AsName(cst.Name("exc")), ), ), ), "code": "try: pass\nexcept* Exception as exc: pass\n", "parser": native_parse_statement, "expected_position": CodeRange((1, 0), (2, 30)), }, # Try/except with multiple clauses { "node": cst.TryStar( cst.SimpleStatementSuite((cst.Pass(), )), handlers=( cst.ExceptStarHandler( cst.SimpleStatementSuite((cst.Pass(), )), type=cst.Name("TypeError"), name=cst.AsName(cst.Name("e")), ), cst.ExceptStarHandler( cst.SimpleStatementSuite((cst.Pass(), )), type=cst.Name("KeyError"), name=cst.AsName(cst.Name("e")), ), ), ), "code": "try: pass\n" + "except* TypeError as e: pass\n" + "except* KeyError as e: pass\n", "parser": native_parse_statement, "expected_position": CodeRange((1, 0), (3, 27)), }, # Simple try/except/finally block { "node": cst.TryStar( cst.SimpleStatementSuite((cst.Pass(), )), handlers=(cst.ExceptStarHandler( cst.SimpleStatementSuite((cst.Pass(), )), type=cst.Name("KeyError"), whitespace_after_except=cst.SimpleWhitespace(""), ), ), finalbody=cst.Finally(cst.SimpleStatementSuite( (cst.Pass(), ))), ), "code": "try: pass\nexcept* KeyError: pass\nfinally: pass\n", "parser": native_parse_statement, "expected_position": CodeRange((1, 0), (3, 13)), }, # Simple try/except/else block { "node": cst.TryStar( cst.SimpleStatementSuite((cst.Pass(), )), handlers=(cst.ExceptStarHandler( cst.SimpleStatementSuite((cst.Pass(), )), type=cst.Name("KeyError"), whitespace_after_except=cst.SimpleWhitespace(""), ), ), orelse=cst.Else(cst.SimpleStatementSuite((cst.Pass(), ))), ), "code": "try: pass\nexcept* KeyError: pass\nelse: pass\n", "parser": native_parse_statement, "expected_position": CodeRange((1, 0), (3, 10)), }, # Verify whitespace in various locations { "node": cst.TryStar( leading_lines=(cst.EmptyLine(comment=cst.Comment("# 1")), ), body=cst.SimpleStatementSuite((cst.Pass(), )), handlers=(cst.ExceptStarHandler( leading_lines=(cst.EmptyLine( comment=cst.Comment("# 2")), ), type=cst.Name("TypeError"), name=cst.AsName( cst.Name("e"), whitespace_before_as=cst.SimpleWhitespace(" "), whitespace_after_as=cst.SimpleWhitespace(" "), ), whitespace_after_except=cst.SimpleWhitespace(" "), whitespace_after_star=cst.SimpleWhitespace(""), whitespace_before_colon=cst.SimpleWhitespace(" "), body=cst.SimpleStatementSuite((cst.Pass(), )), ), ), orelse=cst.Else( leading_lines=(cst.EmptyLine( comment=cst.Comment("# 3")), ), body=cst.SimpleStatementSuite((cst.Pass(), )), whitespace_before_colon=cst.SimpleWhitespace(" "), ), finalbody=cst.Finally( leading_lines=(cst.EmptyLine( comment=cst.Comment("# 4")), ), body=cst.SimpleStatementSuite((cst.Pass(), )), whitespace_before_colon=cst.SimpleWhitespace(" "), ), whitespace_before_colon=cst.SimpleWhitespace(" "), ), "code": "# 1\ntry : pass\n# 2\nexcept *TypeError as e : pass\n# 3\nelse : pass\n# 4\nfinally : pass\n", "parser": native_parse_statement, "expected_position": CodeRange((2, 0), (8, 14)), }, # Now all together { "node": cst.TryStar( cst.SimpleStatementSuite((cst.Pass(), )), handlers=( cst.ExceptStarHandler( cst.SimpleStatementSuite((cst.Pass(), )), type=cst.Name("TypeError"), name=cst.AsName(cst.Name("e")), ), cst.ExceptStarHandler( cst.SimpleStatementSuite((cst.Pass(), )), type=cst.Name("KeyError"), name=cst.AsName(cst.Name("e")), ), ), orelse=cst.Else(cst.SimpleStatementSuite((cst.Pass(), ))), finalbody=cst.Finally(cst.SimpleStatementSuite( (cst.Pass(), ))), ), "code": "try: pass\n" + "except* TypeError as e: pass\n" + "except* KeyError as e: pass\n" + "else: pass\n" + "finally: pass\n", "parser": native_parse_statement, "expected_position": CodeRange((1, 0), (5, 13)), }, )) def test_valid(self, **kwargs: Any) -> None: self.validate_node(**kwargs)
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 WithTest(CSTNodeTest): maxDiff: int = 2000 @data_provider(( # Simple with block { "node": cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), ), "code": "with context_mgr(): pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (1, 24)), }, # Simple async with block { "node": cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), asynchronous=cst.Asynchronous(), ), "code": "async with context_mgr(): pass\n", "parser": lambda code: parse_statement( code, config=PartialParserConfig(python_version="3.7")), }, # Python 3.6 async with block { "node": cst.FunctionDef( cst.Name("foo"), cst.Parameters(), cst.IndentedBlock((cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), asynchronous=cst.Asynchronous(), ), )), asynchronous=cst.Asynchronous(), ), "code": "async def foo():\n async with context_mgr(): pass\n", "parser": lambda code: parse_statement( code, config=PartialParserConfig(python_version="3.6")), }, # Multiple context managers { "node": cst.With( ( cst.WithItem(cst.Call(cst.Name("foo"))), cst.WithItem(cst.Call(cst.Name("bar"))), ), cst.SimpleStatementSuite((cst.Pass(), )), ), "code": "with foo(), bar(): pass\n", "parser": None, }, { "node": cst.With( ( cst.WithItem( cst.Call(cst.Name("foo")), comma=cst.Comma( whitespace_after=cst.SimpleWhitespace(" ")), ), cst.WithItem(cst.Call(cst.Name("bar"))), ), cst.SimpleStatementSuite((cst.Pass(), )), ), "code": "with foo(), bar(): pass\n", "parser": parse_statement, }, # With block containing variable for context manager. { "node": cst.With( (cst.WithItem( cst.Call(cst.Name("context_mgr")), cst.AsName(cst.Name("ctx")), ), ), cst.SimpleStatementSuite((cst.Pass(), )), ), "code": "with context_mgr() as ctx: pass\n", "parser": parse_statement, }, # indentation { "node": DummyIndentedBlock( " ", cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), ), ), "code": " with context_mgr(): pass\n", "parser": None, "expected_position": CodeRange((1, 4), (1, 28)), }, # with an indented body { "node": DummyIndentedBlock( " ", cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.IndentedBlock((cst.SimpleStatementLine( (cst.Pass(), )), )), ), ), "code": " with context_mgr():\n pass\n", "parser": None, "expected_position": CodeRange((1, 4), (2, 12)), }, # leading_lines { "node": cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), leading_lines=(cst.EmptyLine( comment=cst.Comment("# leading comment")), ), ), "code": "# leading comment\nwith context_mgr(): pass\n", "parser": parse_statement, "expected_position": CodeRange((2, 0), (2, 24)), }, # Whitespace { "node": cst.With( (cst.WithItem( cst.Call(cst.Name("context_mgr")), cst.AsName( cst.Name("ctx"), whitespace_before_as=cst.SimpleWhitespace(" "), whitespace_after_as=cst.SimpleWhitespace(" "), ), ), ), cst.SimpleStatementSuite((cst.Pass(), )), whitespace_after_with=cst.SimpleWhitespace(" "), whitespace_before_colon=cst.SimpleWhitespace(" "), ), "code": "with context_mgr() as ctx : pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (1, 36)), }, # Weird spacing rules, that parse differently depending on whether # we are using a grammar that included parenthesized with statements. { "node": cst.With( (cst.WithItem( cst.Call( cst.Name("context_mgr"), lpar=() if is_native() else (cst.LeftParen(), ), rpar=() if is_native() else (cst.RightParen(), ), )), ), cst.SimpleStatementSuite((cst.Pass(), )), lpar=(cst.LeftParen() if is_native() else MaybeSentinel.DEFAULT), rpar=(cst.RightParen() if is_native() else MaybeSentinel.DEFAULT), whitespace_after_with=cst.SimpleWhitespace(""), ), "code": "with(context_mgr()): pass\n", "parser": parse_statement, "expected_position": CodeRange((1, 0), (1, 25)), }, # Multi-line parenthesized with. { "node": cst.With( ( cst.WithItem( cst.Call(cst.Name("foo")), comma=cst. Comma(whitespace_after=cst.ParenthesizedWhitespace( first_line=cst.TrailingWhitespace( whitespace=cst.SimpleWhitespace(value="", ), comment=None, newline=cst.Newline(value=None, ), ), empty_lines=[], indent=True, last_line=cst.SimpleWhitespace(value=" ", ), )), ), cst.WithItem(cst.Call(cst.Name("bar")), comma=cst.Comma()), ), cst.SimpleStatementSuite((cst.Pass(), )), lpar=cst.LeftParen(whitespace_after=cst.SimpleWhitespace(" ")), rpar=cst.RightParen( whitespace_before=cst.SimpleWhitespace(" ")), ), "code": ("with ( foo(),\n" " bar(), ): pass\n"), # noqa "parser": parse_statement if is_native() else None, "expected_position": CodeRange((1, 0), (2, 21)), }, )) def test_valid(self, **kwargs: Any) -> None: self.validate_node(**kwargs) @data_provider(( { "get_node": lambda: cst.With((), cst.IndentedBlock((cst.SimpleStatementLine( (cst.Pass(), )), ))), "expected_re": "A With statement must have at least one WithItem", }, { "get_node": lambda: cst.With( (cst.WithItem( cst.Call(cst.Name("foo")), comma=cst.Comma(whitespace_after=cst.SimpleWhitespace(" ") ), ), ), cst.IndentedBlock((cst.SimpleStatementLine((cst.Pass(), )), )), ), "expected_re": "The last WithItem in an unparenthesized With cannot " + "have a trailing comma.", }, { "get_node": lambda: cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), whitespace_after_with=cst.SimpleWhitespace(""), ), "expected_re": "Must have at least one space after with keyword", }, { "get_node": lambda: cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), whitespace_after_with=cst.SimpleWhitespace(""), lpar=cst.LeftParen(), ), "expected_re": "Do not mix concrete LeftParen/RightParen with " + "MaybeSentinel", }, { "get_node": lambda: cst.With( (cst.WithItem(cst.Call(cst.Name("context_mgr"))), ), cst.SimpleStatementSuite((cst.Pass(), )), whitespace_after_with=cst.SimpleWhitespace(""), rpar=cst.RightParen(), ), "expected_re": "Do not mix concrete LeftParen/RightParen with " + "MaybeSentinel", }, )) def test_invalid(self, **kwargs: Any) -> None: self.assert_invalid(**kwargs) @data_provider(( { "code": "with a, b: pass", "parser": parse_statement_as(python_version="3.1"), "expect_success": True, }, { "code": "with a, b: pass", "parser": parse_statement_as(python_version="3.0"), "expect_success": False, }, )) def test_versions(self, **kwargs: Any) -> None: if is_native() and not kwargs.get("expect_success", True): self.skipTest("parse errors are disabled for native parser") self.assert_parses(**kwargs) def test_adding_parens(self) -> None: node = cst.With( ( cst.WithItem( cst.Call(cst.Name("foo")), comma=cst.Comma( whitespace_after=cst.ParenthesizedWhitespace(), ), ), cst.WithItem(cst.Call(cst.Name("bar")), comma=cst.Comma()), ), cst.SimpleStatementSuite((cst.Pass(), )), lpar=cst.LeftParen(whitespace_after=cst.SimpleWhitespace(" ")), rpar=cst.RightParen(whitespace_before=cst.SimpleWhitespace(" ")), ) module = cst.Module([]) self.assertEqual( module.code_for_node(node), ("with ( foo(),\n" "bar(), ): pass\n") # noqa )