First working version

internal/generator/model.go

@@ -0,0 +1,124 @@
+package generator
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/entity-maker/entity-maker/internal/naming"
+)
+
+// GenerateModel generates the dataclass model
+func GenerateModel(ctx *Context) (string, error) {
+	var b strings.Builder
+
+	// Imports
+	b.WriteString("from dataclasses import dataclass\n")
+
+	// Check what we need to import
+	needsDatetime := NeedsDatetimeImport(ctx.TableInfo.Columns)
+	needsDecimal := NeedsDecimalImport(ctx.TableInfo.Columns)
+
+	if needsDatetime {
+		b.WriteString("from datetime import datetime\n")
+	}
+	if needsDecimal {
+		b.WriteString("from decimal import Decimal\n")
+	}
+
+	b.WriteString("\n")
+	b.WriteString("from televend_core.databases.base_model import Base\n")
+
+	// Import related models for foreign keys
+	fkImports := make(map[string]string) // module_name -> entity_name
+	for _, fk := range ctx.TableInfo.ForeignKeys {
+		moduleName := GetRelationshipModuleName(fk.ForeignTableName)
+		entityName := GetRelationshipEntityName(fk.ForeignTableName)
+		fkImports[moduleName] = entityName
+	}
+
+	// Import enum types
+	if len(ctx.TableInfo.EnumTypes) > 0 {
+		b.WriteString(fmt.Sprintf("from televend_core.databases.televend_repositories.%s.enum import (\n",
+			ctx.ModuleName))
+		for _, enumType := range ctx.TableInfo.EnumTypes {
+			enumName := naming.ToPascalCase(enumType.TypeName)
+			b.WriteString(fmt.Sprintf("    %s,\n", enumName))
+		}
+		b.WriteString(")\n")
+	}
+
+	// Write foreign key imports
+	for moduleName, entityName := range fkImports {
+		b.WriteString(fmt.Sprintf("from televend_core.databases.televend_repositories.%s.model import %s\n",
+			moduleName, entityName))
+	}
+
+	b.WriteString("\n\n")
+
+	// Class definition
+	b.WriteString("@dataclass\n")
+	b.WriteString(fmt.Sprintf("class %s(Base):\n", ctx.EntityName))
+
+	// Get required and optional columns
+	requiredCols := GetRequiredColumns(ctx.TableInfo.Columns)
+	optionalCols := GetOptionalColumns(ctx.TableInfo.Columns)
+
+	// Required fields (non-nullable, no default, not auto-increment, not PK)
+	for _, col := range requiredCols {
+		fieldName := col.Name
+		pythonType := GetPythonTypeForColumn(col, ctx)
+
+		// Regular field
+		b.WriteString(fmt.Sprintf("    %s: %s\n", fieldName, pythonType))
+
+		// Add relationship field if this is a foreign key and column ends with _id
+		// (to avoid name clashes with FK columns that don't follow _id convention)
+		if fk := GetForeignKeyForColumn(col.Name, ctx.TableInfo.ForeignKeys); fk != nil {
+			if strings.HasSuffix(col.Name, "_id") {
+				relationName := GetRelationshipName(col.Name)
+				entityName := GetRelationshipEntityName(fk.ForeignTableName)
+				b.WriteString(fmt.Sprintf("    %s: %s\n", relationName, entityName))
+			}
+		}
+	}
+
+	// Empty line between required and optional
+	if len(requiredCols) > 0 && len(optionalCols) > 0 {
+		b.WriteString("\n")
+	}
+
+	// Optional fields
+	for _, col := range optionalCols {
+		// Skip primary key, we'll add it at the end
+		if col.IsPrimaryKey {
+			continue
+		}
+
+		fieldName := col.Name
+		pythonType := GetPythonTypeForColumn(col, ctx)
+
+		b.WriteString(fmt.Sprintf("    %s: %s | None = None\n", fieldName, pythonType))
+
+		// Add relationship field if this is a foreign key and column ends with _id
+		// (to avoid name clashes with FK columns that don't follow _id convention)
+		if fk := GetForeignKeyForColumn(col.Name, ctx.TableInfo.ForeignKeys); fk != nil {
+			if strings.HasSuffix(col.Name, "_id") {
+				relationName := GetRelationshipName(col.Name)
+				entityName := GetRelationshipEntityName(fk.ForeignTableName)
+				b.WriteString(fmt.Sprintf("    %s: %s | None = None\n", relationName, entityName))
+			}
+		}
+	}
+
+	// Add primary key at the end
+	for _, col := range ctx.TableInfo.Columns {
+		if col.IsPrimaryKey {
+			b.WriteString("\n")
+			pythonType := GetPythonTypeForColumn(col, ctx)
+			b.WriteString(fmt.Sprintf("    %s: %s | None = None\n", col.Name, pythonType))
+			break
+		}
+	}
+
+	return b.String(), nil
+}