mycroforge/MicroForge.CLI/CodeGen/PythonSourceModifier.cs

using Antlr4.Runtime;
using MicroForge.Parsing;

namespace MicroForge.CLI.CodeGen;

public abstract class PythonSourceModifier : PythonParserBaseVisitor<object?>
{
    private CommonTokenStream Stream { get; }
    private PythonParser Parser { get; }
    private TokenStreamRewriter Rewriter { get; }

    protected PythonSourceModifier(string source)
    {
        var input = new AntlrInputStream(source);
        var lexer = new PythonLexer(input);
        Stream = new CommonTokenStream(lexer);
        Parser = new PythonParser(Stream);
        Rewriter = new TokenStreamRewriter(Stream);
    }

    public string Rewrite()
    {
        var tree = Parser.file_input();
        Visit(tree);
        return Rewriter.GetText();
    }

    protected string GetOriginalText(ParserRuleContext context)
    {
        // The parser does not necessarily return the original source,
        // so we return the text from Rewriter.TokenStream, since this is unmodified.
        return Rewriter.TokenStream.GetText(context);
    }

    protected void Rewrite(ParserRuleContext context, params string[] text)
    {
        Rewriter.Replace(from: context.start, to: context.Stop, text: string.Join('\n', text));
    }
}