mycroforge/MycroForge.CLI/CodeGen/OrmEnvUpdater.cs

using MycroForge.Parsing;

namespace MycroForge.CLI.CodeGen;

public class OrmEnvUpdater : PythonSourceModifier
{
    private readonly string _moduleName;
    private readonly string _className;
    private PythonParser.Import_fromContext? _lastImport;

    public OrmEnvUpdater(string source, string moduleName, string className) : base(source)
    {
        _moduleName = moduleName;
        _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 text = GetOriginalText(_lastImport);

        Rewrite(_lastImport, [
            text,
            $"from orm.entities.{_moduleName} import {_className}"
        ]);

        return Rewriter.GetText();
    }

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

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