132 lines
4.6 KiB
C#
132 lines
4.6 KiB
C#
using Humanizer;
|
|
using MycroForge.Parsing;
|
|
|
|
namespace MycroForge.CLI.CodeGen;
|
|
|
|
public partial class EntityLinker
|
|
{
|
|
public class EntityModel : PythonSourceModifier
|
|
{
|
|
private readonly string _className;
|
|
private readonly string _path;
|
|
|
|
private readonly List<PythonParser.Import_fromContext> _importCtxs;
|
|
private readonly List<string> _importsBuffer;
|
|
private PythonParser.Import_fromContext LastImport => _importCtxs.Last();
|
|
|
|
private readonly List<PythonParser.Class_defContext> _classCtxs;
|
|
private PythonParser.AssignmentContext _tableCtx;
|
|
|
|
private readonly List<PythonParser.AssignmentContext> _columnCtxs;
|
|
private readonly List<string> _columnsBuffer;
|
|
private PythonParser.AssignmentContext LastColumn => _columnCtxs.Last();
|
|
|
|
public string ClassName => _className;
|
|
public string Path => _path;
|
|
public string FieldName => _className.Underscore().ToLower();
|
|
|
|
public string TableName => GetOriginalText(_tableCtx)
|
|
.Replace("__tablename__", string.Empty)
|
|
.Replace("=", string.Empty)
|
|
.Replace("\"", string.Empty)
|
|
.Trim();
|
|
|
|
public EntityModel(string className, string path, string source) : base(source)
|
|
{
|
|
_className = className;
|
|
_path = path;
|
|
_importCtxs = new();
|
|
_importsBuffer = new();
|
|
_classCtxs = new();
|
|
_tableCtx = default!;
|
|
_columnCtxs = new();
|
|
_columnsBuffer = new();
|
|
}
|
|
|
|
public void Initialize()
|
|
{
|
|
var tree = Parser.file_input();
|
|
|
|
Visit(tree);
|
|
|
|
if (!_classCtxs.Any(c => GetOriginalText(c).Contains(_className)))
|
|
throw new Exception($"Entity {_className} was not found in {_path}.");
|
|
|
|
if (_columnCtxs.Count == 0)
|
|
throw new Exception($"Entity {_className} has no columns.");
|
|
|
|
_importsBuffer.Add(GetOriginalText(LastImport));
|
|
_columnsBuffer.Add(GetOriginalText(LastColumn));
|
|
|
|
InsertRelationshipImport();
|
|
InsertForeignKeyImport();
|
|
}
|
|
|
|
private void InsertRelationshipImport()
|
|
{
|
|
var relationship = _importCtxs.FirstOrDefault(import =>
|
|
{
|
|
var text = GetOriginalText(import);
|
|
return text.Contains("sqlalchemy.orm") && text.Contains("relationship");
|
|
});
|
|
|
|
if (relationship is null)
|
|
_importsBuffer.Add("from sqlalchemy.orm import relationship");
|
|
}
|
|
|
|
private void InsertForeignKeyImport()
|
|
{
|
|
var foreignKey = _importCtxs.FirstOrDefault(import =>
|
|
{
|
|
var text = GetOriginalText(import);
|
|
return text.Contains("sqlalchemy") && text.Contains("ForeignKey");
|
|
});
|
|
|
|
if (foreignKey is null)
|
|
_importsBuffer.Add("from sqlalchemy import ForeignKey");
|
|
}
|
|
|
|
public override object? VisitImport_from(PythonParser.Import_fromContext context)
|
|
{
|
|
_importCtxs.Add(context);
|
|
return base.VisitImport_from(context);
|
|
}
|
|
|
|
public override object? VisitClass_def(PythonParser.Class_defContext context)
|
|
{
|
|
_classCtxs.Add(context);
|
|
return base.VisitClass_def(context);
|
|
}
|
|
|
|
public override object? VisitAssignment(PythonParser.AssignmentContext context)
|
|
{
|
|
var text = GetOriginalText(context);
|
|
|
|
if (text.StartsWith("__tablename__"))
|
|
_tableCtx = context;
|
|
|
|
if (text.Contains("Mapped["))
|
|
_columnCtxs.Add(context);
|
|
|
|
return base.VisitAssignment(context);
|
|
}
|
|
|
|
public void AppendColumn(params string[] text) => _columnsBuffer.Add(string.Join('\n', text));
|
|
|
|
public void AppendImportIfNotExists(string module, string import, params string[] text)
|
|
{
|
|
var isExisting = _importCtxs.Select(GetOriginalText).Any(ctx => ctx.Contains(module) && ctx.Contains(import));
|
|
var isBuffered = _importsBuffer.Any(txt => txt.Contains(module) && txt.Contains(import));
|
|
|
|
if (!isExisting && !isBuffered)
|
|
_importsBuffer.Add(string.Join('\n', text));
|
|
}
|
|
|
|
public override string Rewrite()
|
|
{
|
|
Rewrite(LastImport, _importsBuffer.ToArray());
|
|
Rewrite(LastColumn, _columnsBuffer.ToArray());
|
|
return Rewriter.GetText();
|
|
}
|
|
}
|
|
} |