optional_typing.py

"""Add Optional typing to Python files."""
import pathlib
import re
from pathlib import Path


def add_optional_typing(file: Path) -> None:
    """Add Optional typing to Python files."""
    file_contents = file.read_text()
    missing_optional_re = re.compile(r'(\w+: )("?\S+"?)( = None)')
    new_file_contents = re.sub(
        missing_optional_re,
        repl=lambda match: f"{match.group(1)}Optional[{match.group(2)}]{match.group(3)}",
        string=file_contents,
    )
    file.write_text(new_file_contents)


def fix_double_option(file: Path) -> None:
    """Fix edge cases where Optional is stacked up."""
    file_contents = file.read_text()
    double_option_re = re.compile(r"(Optional\[Optional\[)(.+)(\]\])")
    new_file_contents = re.sub(
        double_option_re,
        repl=lambda match: f"Optional[{match.group(2)}]",
        string=file_contents,
    )
    file.write_text(new_file_contents)


python_files = pathlib.Path(".").glob("**/*.py")
for file in python_files:
    add_optional_typing(file)
    fix_double_option(file)