mycroforge/MycroForge.CLI/CodeGen/DbEnvUpdater.cs

using MycroForge.Parsing;

namespace MycroForge.CLI.CodeGen;

public class DbEnvUpdater : PythonSourceModifier
{
    private readonly string _importPath;
    private readonly string _className;
    private PythonParser.Import_fromContext? _lastImport;

    public DbEnvUpdater(string source, string importPath, string className) : base(source)
    {
        _importPath = importPath;
        _className = className;
    }

    public override string Rewrite()
    {
        var tree = Parser.file_input();
        Visit(tree);

        if (_lastImport is null)
            throw new Exception("Could not find import insertion point.");
        
        var lastImportText = GetOriginalText(_lastImport);

        Rewrite(_lastImport, [
            lastImportText,
            $"from db.entities.{_importPath} import {_className}"
        ]);

        return Rewriter.GetText();
    }

    public override object? VisitImport_from(PythonParser.Import_fromContext context)
    {
        var text = GetOriginalText(context);

        if (text.StartsWith("from db.entities"))
            _lastImport = context;
        
        return base.VisitImport_from(context);
    }
}