Transform long SimpleStirng of Arg to ConcatenatedString?

[gkze]

Hi, first off -- thanks for this amazing library, I've had a ton of fun using it on my codebase.

I'm trying to address a case where there is a string argument (SimpleString) that exceeds the desired line length:

[...]
    description="Some very long description that exceeds the maximum line length of 80 characters and so on and so forth"
[...]

I'd like to have this be transformed into a parenthesized concatenated string that is split/wrapped such that each of its elements, including the proper indentation, do no exceed 80 characters in line length:

[...]
    description=(
        "Some very long description that exceeds the maximum line length of 80 "
        "characters and so on and so forth"
    )
[...]

Here is as far as I've gotten:

wrap_strings.py

#!/usr/bin/env python3
# To invoke, from $ROOT/api:
# python -m libcst.tool codemod wrap_strings.WrapStringsCommand **/*.py
"""LibCST codemod for wrapping long string arguments in Python files.

This module provides a codemod that identifies strings that exceed a configured
line length (default 80 characters) and wraps them to multiple lines while
preserving indentation. It handles strings in various contexts including
function definitions, if statements, try blocks, and parenthesized expressions.
"""

import argparse
import sys
from collections.abc import Collection
from pathlib import Path
from typing import ClassVar

from libcst import (
    Arg,
    Comment,
    ConcatenatedString,
    CSTTransformer,
    EmptyLine,
    Newline,
    SimpleString,
    SimpleWhitespace,
    matchers,
)
from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
from libcst.metadata import (
    PositionProvider,
    ProviderT,
)
from libcst.tool import main as libcst_main


class SimpleStringToConcatenatedString(CSTTransformer):
    def __init__(self, line_length: int) -> None:
        self.line_length: int = line_length


class WrapStringsCommand(VisitorBasedCodemodCommand):
    DESCRIPTION: str = (
        "Wraps strings that form lines longer than 80 characters."
    )
    LINE_LENGTH: int = 80
    METADATA_DEPENDENCIES: ClassVar[Collection[ProviderT]] = (PositionProvider,)

    @staticmethod
    def parse_args(arg_parser: argparse.ArgumentParser) -> None:
        _ = arg_parser.add_argument(
            "-l",
            "--line-length",
            dest="line_length",
            help="Line length beyond which to wrap strings",
            type=int,
            default=WrapStringsCommand.LINE_LENGTH,
            required=False,
        )

    def __init__(
        self,
        context: CodemodContext,
        line_length: int = LINE_LENGTH,
    ) -> None:
        super().__init__(context)
        self.line_length: int = line_length

    def leave_Arg(
        self,
        original_node: Arg,
        updated_node: Arg,
    ) -> Arg:
        if not matchers.matches(
            updated_node,
            matchers.Arg(
                value=matchers.OneOf(
                    matchers.SimpleString(),
                    matchers.ConcatenatedString(),
                )
            ),
        ):
            return updated_node

        match original_node.value:
            case SimpleString():
                return updated_node.with_changes(
                    value=updated_node.value.visit(
                        SimpleStringToConcatenatedString(self.line_length)
                    )
                )

            case ConcatenatedString():
                return updated_node

        return updated_node


if __name__ == "__main__":
    parser = argparse.ArgumentParser(Path(__file__).name)
    parser.add_argument("-s", "--strip", type=str, default="")
    parser.add_argument("libcst_args", nargs=argparse.REMAINDER)

    filenames = []
    args = parser.parse_args(sys.argv[1:])

    if args.strip:
        for filename in args.libcst_args:
            if filename.strip("/").startswith(args.strip.lstrip("/")):
                path = Path(filename)
                filenames.append(str(path.relative_to(path.parts[0])))
                continue

            filenames.append(filename)

    libcst_main(
        "libcst.tool",
        ["codemod", "wrap_strings.WrapStringsCommand"]
        + (filenames or args.libcst_args),
    )

I've got several questions:

• Is this the right approach? On a few of the previous codemods I've found that I was under-utilizing LibCST APIs and writing a bunch of custom logic instead.
• I've gotten the matchers to filter out unneeded cases, so I've landed on the two that I'd like to address:
  • SimpleString that's longer than 80 chars with indent (using PositionProvider, not sure if there's a better way)
  • ConcatenatedString where any of the elements is longer than 80 chars. In this case I'd like to collect all of the elements and re-produce a correctly concatenated string none of the elements of which are longer than 80 chars with indent.

Conceptually, I've taken a look at how ConcatenatedStrings are expressed, and it seems like I'd have to do some sort of a recursion where I slice off the left side of the string up to line_length - indent index, assign it to the left side of the ConcatenatedString, and pass the remainder down to the same function, until the right side of the string + indent does not exceed 80 chars.

I'm a bit stumped on how to do this because when I tried assigning left and right on an empty ConcatenatedString I realized they're read-only (makes sense), so I probably need to build some sort of a recursive visitor that performs the logic I described above.

What's the best way to move forward?

Thanks in advance for any help.

[gkze]

Quick self-answer here - solved it like this:

wrap_strings.py

#!/usr/bin/env python3
# To invoke, from $ROOT/api:
# python -m libcst.tool codemod wrap_strings.WrapStringsCommand **/*.py
"""LibCST codemod for wrapping long string arguments in Python files.

This module provides a codemod that identifies strings that exceed a configured
line length (default 80 characters) and wraps them to multiple lines while
preserving indentation. It handles strings in various contexts including
function definitions, if statements, try blocks, and parenthesized expressions.
"""

import argparse
import textwrap
from collections.abc import Collection
from typing import ClassVar, cast

from libcst import (
    Arg,
    ConcatenatedString,
    CSTNode,
    CSTTransformer,
    LeftParen,
    Newline,
    ParenthesizedWhitespace,
    RightParen,
    SimpleString,
    SimpleWhitespace,
    TrailingWhitespace,
    matchers,
)
from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
from libcst.metadata import (
    CodeRange,
    ParentNodeProvider,
    PositionProvider,
    ProviderT,
)


class SimpleStringToConcatenatedString(CSTTransformer):
    def __init__(
        self, line_length: int, position: CodeRange, parent: CSTNode
    ) -> None:
        self.line_length: int = line_length
        self.position = position
        self.parent = parent

    def build_concatenated_string(
        self, to_concatenate: list[str], indent: int
    ) -> ConcatenatedString | SimpleString:
        if len(to_concatenate) == 1:
            final_string = to_concatenate[0].strip('"')
            return SimpleString(value=f'"{final_string}"')

        left = to_concatenate.pop(0).strip('"')

        return ConcatenatedString(
            left=SimpleString(value=f'"{left}"'),
            right=self.build_concatenated_string(to_concatenate, indent),
            whitespace_between=ParenthesizedWhitespace(
                first_line=TrailingWhitespace(
                    whitespace=SimpleWhitespace(value=""),
                    comment=None,
                    newline=Newline(value=None),
                ),
                empty_lines=[],
                indent=True,
                last_line=SimpleWhitespace(value=" " * indent),
            ),
        )

    def leave_SimpleString(
        self,
        original_node: SimpleString,
        _: SimpleString,
    ) -> ConcatenatedString | SimpleString:
        res = self.build_concatenated_string(
            textwrap.wrap(
                original_node.value,
                self.line_length - self.position.start.column - 5,
                replace_whitespace=False,
                break_long_words=False,
                drop_whitespace=False,
            ),
            self.position.start.column,
        )

        if not matchers.matches(
            self.parent, matchers.Call(args=[matchers.Arg()])
        ):
            res = res.with_changes(lpar=[LeftParen()], rpar=[RightParen()])

        return res


class WrapStringsCommand(VisitorBasedCodemodCommand):
    DESCRIPTION: str = (
        "Wraps strings that form lines longer than 80 characters."
    )
    LINE_LENGTH: int = 80
    METADATA_DEPENDENCIES: ClassVar[Collection[ProviderT]] = (
        PositionProvider,
        ParentNodeProvider,
    )

    @staticmethod
    def parse_args(arg_parser: argparse.ArgumentParser) -> None:
        _ = arg_parser.add_argument(
            "-l",
            "--line-length",
            dest="line_length",
            help="Line length beyond which to wrap strings",
            type=int,
            default=WrapStringsCommand.LINE_LENGTH,
            required=False,
        )

    def __init__(
        self,
        context: CodemodContext,
        line_length: int = LINE_LENGTH,
    ) -> None:
        super().__init__(context)
        self.line_length: int = line_length

    def leave_Arg(
        self,
        original_node: Arg,
        updated_node: Arg,
    ) -> Arg:
        if not matchers.matches(
            updated_node,
            matchers.Arg(
                value=matchers.OneOf(
                    matchers.SimpleString(),
                    matchers.ConcatenatedString(),
                )
            ),
        ):
            return updated_node

        pos = cast(
            "CodeRange", self.get_metadata(PositionProvider, original_node)
        )

        match original_node.value:
            case SimpleString():
                if (
                    not pos.start.column + len(original_node.value.value) + 1
                    > self.line_length
                ) or original_node.value.value.startswith('"""'):
                    return updated_node

                return updated_node.with_changes(
                    value=updated_node.value.visit(
                        SimpleStringToConcatenatedString(
                            self.line_length,
                            pos,
                            cast(
                                "CSTNode",
                                self.get_metadata(
                                    ParentNodeProvider, original_node
                                ),
                            ),
                        )
                    )
                )

            case ConcatenatedString():
                return updated_node

        return updated_node

Ended up passing the desired case to a CSTTransformer implementation that recursively built up the concatenated string, optionally not parenthesizing it if it was the only call arg. Works for now but would love to know if there's a better way