Convert Schema Table To Python SQLModel classes

GeneratePythonSQLModel.groovy

import com.intellij.database.model.DasTable
import com.intellij.database.util.Case
import com.intellij.database.util.DasUtil

/*
 * Available context bindings:
 *   SELECTION   Iterable<DasObject>
 *   PROJECT     project
 *   FILES       files helper
 */

typeMapping = [
  (~/(?i)int/)                      : "int",
  (~/(?i)float|double|decimal|real/): "float",
  (~/(?i)datetime|timestamp/)       : "datetime",
  (~/(?i)date/)                     : "datetime",
  (~/(?i)time/)                     : "datetime",
  (~/(?i)/)                         : "str"
]

FILES.chooseDirectoryAndSave("Choose directory", "Choose where to store generated files") { dir ->
  SELECTION.filter { it instanceof DasTable }.each { generate(it, dir) }
}

def generate(table, dir) {
  def className = javaName(table.getName(), true)
  def fields = calcFields(table)
  new File(dir, className + ".py").withPrintWriter { out -> generate(out, className, fields,table.getName()) }
}

def generate(out, className, fields,tableName) {
  def pythonIndent="    "
  out.println "\"\"\"Automatically generated by Python SQLModel script\"\"\""
  out.println "from datetime import datetime"
  out.println "from sqlmodel import SQLModel, Field, Relationship"
  out.println "from typing import Optional, List, Any"
  out.println "from sqlalchemy import UniqueConstraint, Column, TEXT, Index"
  out.println "from sqlalchemy.orm import relationship"
  out.println "from pydantic import BaseModel"
  out.println ""
  out.println ""
  out.println "class $className(SQLModel, table=True):"
  out.println "${pythonIndent}__tablename__ = \"$tableName\""
  fields.each() {
    if (it.annos != "") out.println "${pythonIndent}${it.annos}"
    out.println "${pythonIndent}${it.name}: Optional[${it.type}] = Field(${it.defValue}${it.primary}${it.dataLen})"
  }
  out.println ""

}

def calcFields(table) {
  DasUtil.getColumns(table).reduce([]) { fields, col ->
    def spec = Case.LOWER.apply(col.getDataType().getSpecification())
    def typeStr = typeMapping.find { p, t -> p.matcher(spec).find() }.value
    fields += [[
                 name : Case.LOWER.apply(col.getName()),
                 type : typeStr,
                 annos: "",
                 comment: col.getComment()?", description=\"${col.getComment()}\"":"",
                 defValue: !col.isNotNull()&&!col.getDefault()?"default=None":(col.isNotNull()&&col.getDefault()?"default=\"${col.getDefault()}\"":""),
                 dataLen: col.getDataType().getLength()?", max_length=${col.getDataType().getLength()}":"",
                 maxLen: col.getDataType().size,
                 dataType: col.getDataType().typeName,
                 primary: DasUtil.isPrimary(col)?", primary_key=True":"",
                 autoGenerated: DasUtil.isAutoGenerated(col)
                 ]]
  }
}

def javaName(str, capitalize) {
  def s = com.intellij.psi.codeStyle.NameUtil.splitNameIntoWords(str)
    .collect { Case.LOWER.apply(it).capitalize() }
    .join("")
    .replaceAll(/[^\p{javaJavaIdentifierPart}[_]]/, "_")
  capitalize || s.length() == 1? s : Case.LOWER.apply(s[0]) + s[1..-1]
}

def convertPythonBool(b){
    if(b){
        return "True"
    } else {
        return "False"
       }
}