import os
import re
import sys
from pathlib import Path
from tempfile import NamedTemporaryFile
from typing import Any, TextIO, Union
type AnyIO = Union[TextIO, Any]
class SourceModifier:
def __init__(self):
self.general = GeneralContext()
self.reference = ReferenceContext()
@property
def dirty(self) -> bool:
return self.general.dirty or self.reference.dirty
def modify_file(self, path: Path):
with open(path, 'r') as reader:
with NamedTemporaryFile('w', dir=path.parent) as writer:
self.modify_stream(reader, writer)
if self.dirty:
writer._closer.delete = False
if self.dirty:
print(path)
os.rename(writer.name, path)
def modify_stream(self, reader: TextIO, writer: AnyIO):
for line in reader:
line = line.rstrip('\n')
line = self.general.modify_line(line)
line = self.reference.modify_line(line)
print(line, file=writer)
class GeneralContext:
def __init__(self):
self.dirty = False
def modify_line(self, line: str) -> str:
old = line
line = re.sub(r' +$', '', line)
if match := re.search(r'^#\[derive\((.+)\)]$', line):
tokens = match.group(1).split(',')
tokens = ', '.join(sorted(x.strip() for x in tokens))
line = f'#[derive({tokens})]'
if line != old:
self.dirty = True
return line
class ReferenceContext:
def __init__(self):
self.open = re.compile(r'^(\s*)impl\b.+\b(\w+)\b(<.+?>)? (?:{|where)?$')
self.modify = None
self.close = None
self.dirty = False
def modify_line(self, line: str) -> str:
if match := self.open.search(line):
indent = match.group(1)
block = match.group(2)
template = match.group(3) or ''
self.modify = re.compile(rf'([^:])\b{block}\b({template})?')
self.close = re.compile(rf'^{indent}}}$')
elif self.modify:
old = line
line, quotes, comment = self.remove_quotes(line)
line = self.modify.sub(r'\1Self', line)
line = self.insert_quotes(line, quotes, comment)
if line != old:
self.dirty = True
if self.close and self.close.search(line):
self.modify = None
self.close = None
return line
@staticmethod
def remove_quotes(line: str) -> tuple[str, list[str], str]:
if match := re.search(r'^(.*?)(\s*//.*)$', line):
line, comment = match.groups()
else:
comment = ''
quotes = re.findall(r'"(.*?)"', line)
if quotes:
line = re.sub(r'"(.*?)"', '""', line)
return line, quotes, comment
@staticmethod
def insert_quotes(line: str, quotes: list[str], comment: str) -> str:
if quotes:
index = 0
def inner(_: re.Match) -> str:
nonlocal index
quote = quotes[index] if index < len(quotes) else ''
index += 1
return f'"{quote}"'
line = re.sub(r'""', inner, line)
return line + comment
def run_main():
_, *paths = sys.argv
if paths:
for path in paths:
path = Path(path)
if path.is_dir():
for rootdir, subdir, files in path.walk():
for file in files:
modifier = SourceModifier()
modifier.modify_file(rootdir / file)
elif path.is_file():
modifier = SourceModifier()
modifier.modify_file(path)
else:
modifier = SourceModifier()
modifier.modify_stream(sys.stdin, sys.stdout)
try:
run_main()
except (OSError, KeyboardInterrupt) as error:
print(error, file=sys.stderr)