using Humanizer; namespace MycroForge.CLI.CodeGen; public partial class EntityLinker { public class EntityModel { private readonly Source _source; private readonly string _className; private readonly string _path; private readonly List _importContexts; private readonly List _importsBuffer; private readonly List _classContexts; private SourceMatch _tableContext; private readonly List _columnContexts; private readonly List _columnsBuffer; private SourceMatch LastColumn => _columnContexts.Last(); private SourceMatch LastImport => _importContexts.Last(); public string ClassName => _className; public string Path => _path; public string FieldName => _className.Underscore().ToLower(); public string TableName => _tableContext.Text .Replace("__tablename__", string.Empty) .Replace("=", string.Empty) .Replace("\"", string.Empty) .Trim(); public EntityModel(string className, string path, string source) { _source = new Source(source); _className = className; _path = path; _importContexts = new(); _importsBuffer = new(); _classContexts = new(); _tableContext = default!; _columnContexts = new(); _columnsBuffer = new(); var classContexts = _source.FindAll(@"class\s+.*\s*\(EntityBase\)\s*:\s"); if (!classContexts.Any(c => c.Text.Contains(_className))) throw new Exception($"Entity {_className} was not found in {_path}."); } public void Initialize() { ReadImports(); AssertClassExists(); AssertTableNameExists(); ReadColumns(); InsertRelationshipImport(); InsertForeignKeyImport(); } private void ReadImports() { _importContexts.AddRange(_source.FindAll(@"from\s+.+\s+import\s+.*\s")); } private void AssertClassExists() { _classContexts.AddRange(_source.FindAll(@"class\s+.+\s*\(EntityBase\)\s*:\s")); if (!_classContexts.Any(c => c.Text.Contains(_className))) throw new Exception($"Entity {_className} was not found in {_path}."); } private void AssertTableNameExists() { _tableContext = _source.Find(@"__tablename__\s*=\s*.+\s"); if (string.IsNullOrEmpty(_tableContext.Text)) throw new Exception($"__tablename__ definition for Entity {_className} was not found in {_path}."); } private void ReadColumns() { _columnContexts.AddRange( _source.FindAll(@".+\s*:\s*Mapped\[.+\]\s*=\s*(relationship|mapped_column)\(.+\)\s") ); if (_columnContexts.Count == 0) throw new Exception($"Entity {_className} has no columns."); } private void InsertRelationshipImport() { var relationship = _importContexts.FirstOrDefault( import => import.Text.Contains("sqlalchemy.orm") && import.Text.Contains("relationship") ); if (relationship is null) _importsBuffer.Add("from sqlalchemy.orm import relationship"); } private void InsertForeignKeyImport() { var foreignKey = _importContexts.FirstOrDefault( import => import.Text.Contains("sqlalchemy") && import.Text.Contains("ForeignKey") ); if (foreignKey is null) _importsBuffer.Add("from sqlalchemy import ForeignKey"); } public void AppendColumn(string text) { _columnsBuffer.Add(text); } public void AppendColumns(params string[] text) { _columnsBuffer.Add(string.Join('\n', text)); } public void Import(string from, string import) { var exists = _importContexts.Any(ctx => ctx.Text.Contains(from) && ctx.Text.Contains(import)); var buffered = _importsBuffer.Any(txt => txt.Contains(from) && txt.Contains(import)); if (!exists && !buffered) _importsBuffer.Add($"from {from} import {import}"); } public string Rewrite() { // The order matters! We must first rewrite the columns first, // so that their indexes are not modified when inserting imports. _source.InsertMultiLine(LastColumn.EndIndex, _columnsBuffer.Append("\n").ToArray()); _source.InsertMultiLine(LastImport.EndIndex, _importsBuffer.Append("\n").ToArray()); return _source.Text; } } }