Clean Python Code

Python and PySpark formatting for VS Code with vertical method chains, hierarchical indentation and optional leading commas.

Version 0.4.3: syntax-aware formatting with f-string support

The formatting engine now uses Python's built-in ast and tokenize modules. Operators are recognized from syntax nodes and tokens, so color, cursor, generator expressions and strings such as "pending or review" retain their meaning.

Before applying an edit, the extension checks:

1. The original and formatted abstract syntax trees are equivalent, excluding source locations.
2. All non-parenthesis tokens remain identical, including strings, comments and tuple commas. Parentheses added for line continuation must also preserve the AST.
3. Formatting a second time produces exactly the same result.
4. The surrounding Python document remains equivalent and has not changed during the request.

Code is parsed, never imported or executed by the formatter. No pip packages or Spark installation are needed.

Changes in 0.4.3

• F-strings no longer cause the entire statement to be skipped. Their exact source text is kept opaque, including interpolation expressions, debug whitespace, conversions and format specifications. The surrounding Python/PySpark code is formatted normally.
• Nested f-strings, reused quotes, raw prefixes and multiline f-strings are supported when the selected Python interpreter supports their syntax. Multiline plain-string contents are also preserved while surrounding calls are formatted.
• Return pipelines use one consistent continuation level for the base expression and following methods. Their closing parenthesis aligns with return.
• Starred comprehensions in expanded argument lists receive a readable vertical layout.
• External trailing comments remain attached to the same statement instead of preventing all formatting of that statement.
• The full join example has a reviewed output fixture and a real bridge test. Literal spelling is checked character for character in addition to AST/token equivalence and idempotence.

Changes in 0.4.2

• Expanded list/set/dict comprehensions and generator expressions put the element, for/async for clauses, filters and closing delimiter on separate lines. Short simple comprehensions remain inline.
• Nested fluent calls use the nesting level of their first call, preventing indentation drift across repeated .when() calls. Top-level pipelines retain the existing continuation style.
• A call expands when one of its arguments is multiline, including .otherwise(F.array_union(...)).
• Short predicate suffixes such as .isin(...) remain compact when their arguments are inline.
• Call parentheses are distinguished from outer expression grouping, preventing a nested argument chain from receiving top-level indentation.
• Both full user examples have reviewed before/after fixtures. The release adds 114 combinations of comprehensions and nested chains across three settings profiles.

Changes in 0.4.1

• Multiline indentation is measured from the rendered opening delimiter. Closing delimiters align with that exact opener column.
• Each chained method is measured independently, so a long preceding pipeline does not force .where(), .select() or .distinct() to expand.
• Internal comments no longer cause an entire collection to be skipped. Comment-bearing delimiters use trailing commas to keep section and inline comments attached to their code.
• argumentIndent and outerIndent default to 2 spaces; chainIndent defaults to 0 so continuation dots align with the first chain dot.
• The full reported heatmap example and its reviewed expected output are included in test/fixtures/, along with dedicated visual layout regressions.

Requirements

• VS Code 1.90 or newer, with a trusted workspace.
• Python 3.9 or newer available on the extension host. The interpreter must support the syntax in your document; for example, match requires Python 3.10+ and PEP 701 f-strings require Python 3.12+.
• In SSH, WSL or containers, Python must be available in that remote environment.

Interpreter discovery uses cleanPythonCode.pythonPath first, then the environment selected by the Microsoft Python extension, then python3 / python on PATH (py -3 / python / python3 on Windows). The Microsoft Python extension is optional. An explicitly configured invalid executable produces an error rather than silently using another interpreter.

Use

Select complete Python statements and press Ctrl + Alt + S on Windows, Linux or macOS, or run Clean Python Code: Format Selection from the Command Palette. Python cells in .ipynb notebooks use the same command.

If Python cannot be found, set the executable path in Settings → Clean Python Code → Python Path. Supply an executable path, not a shell command or command-line arguments.

Invalid Python, incomplete selections, unsupported syntax and failed validation leave the original code intact and display a reason. Documents/cells containing IPython magics such as %time or !pip are not ordinary Python and are not supported by this version. Fix syntax errors elsewhere in the document before formatting a selection.

Formatting scope

• Assignments, annotated/augmented assignments, expression statements, returns and raises.
• Calls with multiple arguments, lists, dictionaries and long collections.
• PySpark chains such as .filter(), .select(), .withColumn(), .join() and Window expressions.
• Long Python boolean and bitwise expressions, with their original precedence.
• Nested statements inside functions, classes, loops, conditionals and exception handlers.

Structural headers, import statements, inline suites, semicolon statements and tab-indented statements remain unchanged. Lists, dictionaries and calls containing comments can now be formatted. Comment text, order and inline attachment are preserved; delimiters containing comments retain trailing separators even when leadingComma is enabled. External trailing comments stay on the closing line of the same statement. String and f-string contents are preserved verbatim while surrounding calls, collections and chains can be formatted. Expressions inside f-string interpolation are intentionally not reformatted. Comprehensions and lambdas retain their syntax and can contain formatted nested calls. Comprehensions expand by AST-identified clauses when long or when a nested expression requires multiple lines. Long literal strings are not split. maxInlineLength is a layout preference, not a strict line-length limit.

The formatter keeps existing explicit grouping and comma tokens. It does not add parentheses around every short boolean operand or rewrite expressions using ast.unparse.

Example

Before:

df = spark.table("x").filter(F.col("a") > 0).select("a", "b")

After, with default settings:

df = (
        spark.table("x")
             .filter(F.col("a") > 0)
             .select(
                 "a"
                 , "b"
             )
)

These expressions also retain their original meaning:

out = process(color, other)
n = build(cursor, limit)
total = sum(item.price * item.qty for item in cart)
any_ok = any(check(r) for r in rows)
df = df.filter(F.col("status") == "pending or review")

Settings

Install and publish

Use Extensions → … → Install from VSIX… and select clean-python-code-0.4.3.vsix, then reload VS Code if prompted.

The extension identity remains cleanpythoncode.clean-python-code.

The VSIX includes the implementation, tests, an offline packaging script and REVIEW-0.4.3.md. From its extracted extension directory:

npm test
python3 package_vsix.py --output clean-python-code-0.4.3.vsix

npm test uses Node's built-in test runner (Node 18+ for development) and Python on PATH. Set CLEAN_PYTHON_TEST_EXECUTABLE to choose the test interpreter. Runtime has no npm dependencies. See PUBLISHING.md for release steps and validation limits.

• Safely normalizes problematic Unicode whitespace copied into Python code (for example U+00A0), while preserving strings and comments.

Notebook shortcut behavior

Ctrl + Alt + S formats the current Python selection. If there is no text selection, it formats the active Python notebook cell (or the active Python document). The Output panel Clean Python Code records targeting/formatting diagnostics.

Reviewed fluent-chain layout

The formatter keeps the first fluent segment on the base-expression line and aligns later segments with that first dot:

spark = (
          DatabricksSession.builder
                            .getOrCreate()
        )

df_ai_requests = (
                   spark.read
                        .csv(
                              path,
                              header = True,
                            )
                 )

For Python notebooks, the command also falls back to notebook kernel/language metadata. A code cell incorrectly persisted by VS Code as javascript can still be formatted when the notebook itself is Python; after successful Python validation, the extension attempts to set the cell language mode back to python.