# Copyright (c) Meta Platforms, Inc. and affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. # import argparse import ast from typing import Generator, List, Optional, Sequence, Set, Tuple import libcst as cst import libcst.matchers as m from libcst.codemod import ( CodemodContext, ContextAwareTransformer, ContextAwareVisitor, VisitorBasedCodemodCommand, ) def _get_lhs(field: cst.BaseExpression) -> cst.BaseExpression: if isinstance(field, (cst.Name, cst.Integer)): return field elif isinstance(field, (cst.Attribute, cst.Subscript)): return _get_lhs(field.value) else: raise Exception("Unsupported node type!") def _find_expr_from_field_name( fieldname: str, args: Sequence[cst.Arg] ) -> Optional[cst.BaseExpression]: # Things like "0.name" are invalid expressions in python since # we can't tell if name is supposed to be the fraction or a name. # So we do a trick to parse here where we wrap the LHS in parens # and assume LibCST will handle it. if "." in fieldname: ind, exp = fieldname.split(".", 1) fieldname = f"({ind}).{exp}" field_expr = cst.parse_expression(fieldname) lhs = _get_lhs(field_expr) # Verify we don't have any *args or **kwargs attributes. if any(arg.star != "" for arg in args): return None # Get the index into the arg index: Optional[int] = None if isinstance(lhs, cst.Integer): index = int(lhs.value) if index < 0 or index >= len(args): raise Exception(f"Logic error, arg sequence {index} out of bounds!") elif isinstance(lhs, cst.Name): for i, arg in enumerate(args): kw = arg.keyword if kw is None: continue if kw.value == lhs.value: index = i break if index is None: raise Exception(f"Logic error, arg name {lhs.value} out of bounds!") if index is None: raise Exception(f"Logic error, unsupported fieldname expression {fieldname}!") # Format it! return field_expr.deep_replace(lhs, args[index].value) def _get_field(formatstr: str) -> Tuple[str, Optional[str], Optional[str]]: in_index: int = 0 format_spec: Optional[str] = None conversion: Optional[str] = None # Grab any format spec as long as its not an array slice for pos, char in enumerate(formatstr): if char == "[": in_index += 1 elif char == "]": in_index -= 1 elif char == ":": if in_index == 0: formatstr, format_spec = (formatstr[:pos], formatstr[pos + 1 :]) break # Grab any conversion if "!" in formatstr: formatstr, conversion = formatstr.split("!", 1) # Return it return formatstr, format_spec, conversion def _get_tokens( # noqa: C901 string: str, ) -> Generator[Tuple[str, Optional[str], Optional[str], Optional[str]], None, None]: length = len(string) prefix: str = "" format_accum: str = "" in_brackets: int = 0 seen_escape: bool = False for pos, char in enumerate(string): if seen_escape: # The last character was an escape character, so consume # this one as well, and then pop out of the escape. if in_brackets == 0: prefix += char else: format_accum += char seen_escape = False continue # We can't escape inside a f-string/format specifier. if in_brackets == 0: # Grab the next character to see if we are an escape sequence. next_char: Optional[str] = None if pos < length - 1: next_char = string[pos + 1] # If this current character is an escape, we want to # not react to it, append it to the current accumulator and # then do the same for the next character. if char == "{" and next_char == "{": seen_escape = True if char == "}" and next_char == "}": seen_escape = True # Only if we are not an escape sequence do we consider these # brackets. if not seen_escape: if char == "{": in_brackets += 1 # We want to add brackets to the format accumulator as # long as they aren't the outermost, because format # specs allow {} expansion. if in_brackets == 1: continue if char == "}": in_brackets -= 1 if in_brackets < 0: raise Exception("Stray } in format string!") if in_brackets == 0: field_name, format_spec, conversion = _get_field(format_accum) yield (prefix, field_name, format_spec, conversion) prefix = "" format_accum = "" continue # Place in the correct accumulator if in_brackets == 0: prefix += char else: format_accum += char if in_brackets > 0: raise Exception("Stray { in format string!") if format_accum: raise Exception("Logic error!") # Yield the last bit of information yield (prefix, None, None, None) class StringQuoteGatherer(ContextAwareVisitor): def __init__(self, context: CodemodContext) -> None: super().__init__(context) self.stringends: Set[str] = set() def visit_SimpleString(self, node: cst.SimpleString) -> None: self.stringends.add(node.value[-1]) class StripNewlinesTransformer(ContextAwareTransformer): def leave_ParenthesizedWhitespace( self, original_node: cst.ParenthesizedWhitespace, updated_node: cst.ParenthesizedWhitespace, ) -> cst.SimpleWhitespace: return cst.SimpleWhitespace(" ") class SwitchStringQuotesTransformer(ContextAwareTransformer): def __init__(self, context: CodemodContext, avoid_quote: str) -> None: super().__init__(context) if avoid_quote not in {'"', "'"}: raise Exception("Must specify either ' or \" single quote to avoid.") self.avoid_quote: str = avoid_quote self.replace_quote: str = '"' if avoid_quote == "'" else "'" def leave_SimpleString( self, original_node: cst.SimpleString, updated_node: cst.SimpleString ) -> cst.SimpleString: if self.avoid_quote in updated_node.quote: # Attempt to swap the value out, verify that the string is still identical # before and after transformation. new_quote = updated_node.quote.replace(self.avoid_quote, self.replace_quote) new_value = ( f"{updated_node.prefix}{new_quote}{updated_node.raw_value}{new_quote}" ) try: new_str = ast.literal_eval(new_value) if updated_node.evaluated_value != new_str: # This isn't the same! return updated_node return updated_node.with_changes(value=new_value) except Exception: # Failed to parse string, changing the quoting screwed us up. pass # Either failed to parse the new string, or don't need to make changes. return updated_node class ConvertFormatStringCommand(VisitorBasedCodemodCommand): DESCRIPTION: str = "Converts instances of str.format() to f-string." @staticmethod def add_args(arg_parser: argparse.ArgumentParser) -> None: arg_parser.add_argument( "--allow-strip-comments", dest="allow_strip_comments", help=( "Allow stripping comments inside .format() calls when converting " + "to f-strings." ), action="store_true", ) arg_parser.add_argument( "--allow-await", dest="allow_await", help=( "Allow converting expressions inside .format() calls that contain " + "an await expression (only compatible with Python 3.7+)." ), action="store_true", ) def __init__( self, context: CodemodContext, allow_strip_comments: bool = False, allow_await: bool = False, ) -> None: super().__init__(context) self.allow_strip_comments = allow_strip_comments self.allow_await = allow_await def leave_Call( # noqa: C901 self, original_node: cst.Call, updated_node: cst.Call ) -> cst.BaseExpression: # Lets figure out if this is a "".format() call extraction = self.extract( updated_node, m.Call( func=m.Attribute( value=m.SaveMatchedNode(m.SimpleString(), "string"), attr=m.Name("format"), ) ), ) if extraction is not None: fstring: List[cst.BaseFormattedStringContent] = [] inserted_sequence: int = 0 stringnode = cst.ensure_type(extraction["string"], cst.SimpleString) tokens = _get_tokens(stringnode.raw_value) for literal_text, field_name, format_spec, conversion in tokens: if literal_text: fstring.append(cst.FormattedStringText(literal_text)) if field_name is None: # This is not a format-specification continue # Auto-insert field sequence if it is empty if field_name == "": field_name = str(inserted_sequence) inserted_sequence += 1 # Now, if there is a valid format spec, parse it as a f-string # as well, since it allows for insertion of parameters just # like regular f-strings. format_spec_parts: List[cst.BaseFormattedStringContent] = [] if format_spec is not None and len(format_spec) > 0: # Parse the format spec out as a series of tokens as well. format_spec_tokens = _get_tokens(format_spec) for ( spec_literal_text, spec_field_name, spec_format_spec, spec_conversion, ) in format_spec_tokens: if spec_format_spec is not None: # This shouldn't be possible, we don't allow it in the spec! raise Exception("Logic error!") if spec_literal_text: format_spec_parts.append( cst.FormattedStringText(spec_literal_text) ) if spec_field_name is None: # This is not a format-specification continue # Auto-insert field sequence if it is empty if spec_field_name == "": spec_field_name = str(inserted_sequence) inserted_sequence += 1 # Now, convert the spec expression itself. fstring_expression = self._convert_token_to_fstring_expression( spec_field_name, spec_conversion, updated_node.args, stringnode, ) if fstring_expression is None: return updated_node format_spec_parts.append(fstring_expression) # Finally, output the converted value. fstring_expression = self._convert_token_to_fstring_expression( field_name, conversion, updated_node.args, stringnode ) if fstring_expression is None: return updated_node # Technically its valid to add the parts even if it is empty, but # it results in an empty format spec being added which is ugly. if format_spec_parts: fstring_expression = fstring_expression.with_changes( format_spec=format_spec_parts ) fstring.append(fstring_expression) # We converted each part, so lets bang together the f-string itself. return cst.FormattedString( parts=fstring, start=f"f{stringnode.prefix}{stringnode.quote}", end=stringnode.quote, ) return updated_node def _convert_token_to_fstring_expression( self, field_name: str, conversion: Optional[str], arguments: Sequence[cst.Arg], containing_string: cst.SimpleString, ) -> Optional[cst.FormattedStringExpression]: expr = _find_expr_from_field_name(field_name, arguments) if expr is None: # Most likely they used * expansion in a format. self.warn(f"Unsupported field_name {field_name} in format() call") return None # Verify that we don't have any comments or newlines. Comments aren't # allowed in f-strings, and newlines need parenthesization. We can # have formattedstrings inside other formattedstrings, but I chose not # to doeal with that for now. if self.findall(expr, m.Comment()) and not self.allow_strip_comments: # We could strip comments, but this is a formatting change so # we choose not to for now. self.warn("Unsupported comment in format() call") return None if self.findall(expr, m.FormattedString()): self.warn("Unsupported f-string in format() call") return None if self.findall(expr, m.Await()) and not self.allow_await: # This is fixed in 3.7 but we don't currently have a flag # to enable/disable it. self.warn("Unsupported await in format() call") return None # Stripping newlines is effectively a format-only change. expr = cst.ensure_type( expr.visit(StripNewlinesTransformer(self.context)), cst.BaseExpression, ) # Try our best to swap quotes on any strings that won't fit expr = cst.ensure_type( expr.visit( SwitchStringQuotesTransformer(self.context, containing_string.quote[0]) ), cst.BaseExpression, ) # Verify that the resulting expression doesn't have a backslash # in it. raw_expr_string = self.module.code_for_node(expr) if "\\" in raw_expr_string: self.warn("Unsupported backslash in format expression") return None # For safety sake, if this is a dict/set or dict/set comprehension, # wrap it in parens so that it doesn't accidentally create an # escape. if (raw_expr_string.startswith("{") or raw_expr_string.endswith("}")) and ( not expr.lpar or not expr.rpar ): expr = expr.with_changes(lpar=[cst.LeftParen()], rpar=[cst.RightParen()]) # Verify that any strings we insert don't have the same quote quote_gatherer = StringQuoteGatherer(self.context) expr.visit(quote_gatherer) for stringend in quote_gatherer.stringends: if stringend in containing_string.quote: self.warn("Cannot embed string with same quote from format() call") return None return cst.FormattedStringExpression(expression=expr, conversion=conversion)