Recursive Rewrite Type Has Fix

• Playground Link

Description

Suppose we want to transform Python's Union[T1, T2] to T1 | T2 and Optional[T] to T | None.

By default, ast-grep will only fix the outermost node that matches a pattern and will not rewrite the inner AST nodes inside a match. This avoids unexpected rewriting or infinite rewriting loop.

So if you are using non-recurisve rewriter like this, Optional[Union[int, str]] will only be converted to Union[int, str] | None. Note the inner Union[int, str] is not enabled. This is because the rewriter optional matches Optional[$TYPE] and rewrite it to $TYPE | None. The inner $TYPE is not processed.

However, we can apply rewriters to inner types recursively. Take the optional rewriter as an example, we need to apply rewriters, optional and unioins, recursively to $TYPE and get a new variable $NT.

YAML

id: recursive-rewrite-types
language: python
rewriters:
# rewrite Optional[T] to T | None
- id: optional
  rule:
    any:
    - pattern:
        context: 'arg: Optional[$TYPE]'
        selector: generic_type
    - pattern: Optional[$TYPE]
  # recursively apply rewriters to $TYPE
  transform:
    NT:
      rewrite:
        rewriters: [optional, unions]
        source: $TYPE
  # use the new variable $NT
  fix: $NT | None

# similar to Optional, rewrite Union[T1, T2] to T1 | T2
- id: unions
  language: Python
  rule:
    pattern:
      context: 'a: Union[$$$TYPES]'
      selector: generic_type
  transform:
    UNIONS:
      # rewrite all types inside $$$TYPES
      rewrite:
        rewriters: [ rewrite-unions ]
        source: $$$TYPES
        joinBy: " | "
  fix: $UNIONS
- id: rewrite-unions
  rule:
    pattern: $TYPE
    kind: type
  # recursive part
  transform:
    NT:
      rewrite:
        rewriters: [optional, unions]
        source: $TYPE
  fix: $NT

# find all types
rule:
  kind: type
  pattern: $TPE
# apply the recursive rewriters
transform:
  NEW_TYPE:
    rewrite:
      rewriters: [optional, unions]
      source: $TPE
# output
fix: $NEW_TYPE

Example

results:  Optional[Union[List[Union[str, dict]], str]]

Diff

results:  Optional[Union[List[Union[str, dict]], str]] 
results:  List[str | dict] | str | None

Contributed by

Inspired by steinuil