initial commit
This commit is contained in:
211
internal/helper.go
Normal file
211
internal/helper.go
Normal file
@@ -0,0 +1,211 @@
|
||||
package internal
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type SVGData struct {
|
||||
Name string
|
||||
SVG string
|
||||
}
|
||||
|
||||
func GetSvgsFromPath(path string) ([]SVGData, error) {
|
||||
files, err := os.ReadDir(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
svgs := make([]SVGData, 0)
|
||||
for _, file := range files {
|
||||
if file.IsDir() || !strings.HasSuffix(file.Name(), ".svg") {
|
||||
continue
|
||||
}
|
||||
svgFile, err := os.ReadFile(filepath.Join(path, file.Name()))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
svg := SVGData{
|
||||
Name: file.Name(),
|
||||
SVG: string(svgFile),
|
||||
}
|
||||
svgs = append(svgs, svg)
|
||||
}
|
||||
return svgs, nil
|
||||
}
|
||||
|
||||
type Response struct {
|
||||
Data []struct {
|
||||
URL string `json:"url"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
func GenerateIcon(prompt string) (string, error) {
|
||||
client := &http.Client{}
|
||||
recraftBody := fmt.Sprintf(`{
|
||||
"model": "recraftv2",
|
||||
"prompt": "%s",
|
||||
"style": "icon"
|
||||
}`, prompt)
|
||||
req, err := http.NewRequest("POST", "https://external.api.recraft.ai/v1/images/generations", strings.NewReader(recraftBody))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req.Header.Add("Content-Type", "application/json")
|
||||
req.Header.Add("Authorization", "Bearer i153LYqEDGOipN2LCCvRBCV2uwJWMH7nZzxX8RKiZwfVTGgAgXBMDVNgvEYk0YtJ")
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("bad status: %s", resp.Status)
|
||||
}
|
||||
// Step 2: Read the response body into memory
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
var data Response
|
||||
if err := json.Unmarshal(body, &data); err != nil {
|
||||
return "", fmt.Errorf("unable to unmarshal json: %v", err)
|
||||
}
|
||||
if len(data.Data) == 0 {
|
||||
return "", fmt.Errorf("no data found in response")
|
||||
}
|
||||
return data.Data[0].URL, nil
|
||||
}
|
||||
|
||||
func DownloadFile(filepath string, url string) error {
|
||||
// out, err := os.Create(filepath)
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
// defer out.Close()
|
||||
resp, err := http.Get(url)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("bad status: %s", resp.Status)
|
||||
}
|
||||
svgData, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cleanSvg := CleanHeader(string(svgData))
|
||||
if err = os.WriteFile(filepath, []byte(cleanSvg), 0700); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
const (
|
||||
lowercase = "abcdefghijklmnopqrstuvwxyz"
|
||||
uppercase = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
|
||||
digits = "0123456789"
|
||||
|
||||
allChars = lowercase + uppercase + digits + "-"
|
||||
)
|
||||
|
||||
var Allowed map[rune]struct{} = func() map[rune]struct{} {
|
||||
allowed := make(map[rune]struct{})
|
||||
for _, ch := range allChars {
|
||||
allowed[ch] = struct{}{}
|
||||
}
|
||||
return allowed
|
||||
}()
|
||||
|
||||
func SanatizePrompt(prompt string) string {
|
||||
var sb strings.Builder
|
||||
// Replace disallowed characters with '-'
|
||||
for _, ch := range prompt {
|
||||
if _, ok := Allowed[ch]; ok {
|
||||
sb.WriteRune(ch)
|
||||
} else {
|
||||
sb.WriteRune('-')
|
||||
}
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
func FixIconHeader() {
|
||||
cliCmd := flag.NewFlagSet("fix", flag.ExitOnError)
|
||||
svgPath := cliCmd.String("svg-path", "", "Path to folder with svgs")
|
||||
outputPath := cliCmd.String("output-path", "", "Path to put updated svgs")
|
||||
if err := cliCmd.Parse(os.Args[2:]); err != nil {
|
||||
log.Fatalf("Failed to parse flags: %v", err)
|
||||
}
|
||||
svgs, err := GetSvgsFromPath(*svgPath)
|
||||
if err != nil {
|
||||
log.Fatalf("fatal error reading svg dir %v", err)
|
||||
}
|
||||
if err = os.MkdirAll(*outputPath, 0700); err != nil {
|
||||
log.Fatalf("fatal error making output dir %v", err)
|
||||
}
|
||||
for _, svg := range svgs {
|
||||
cleanSvg := CleanHeader(svg.SVG)
|
||||
if err = os.WriteFile(filepath.Join(*outputPath, svg.Name), []byte(cleanSvg), 0700); err != nil {
|
||||
log.Printf("error writing %s; %v ", svg.Name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func FileExists(path string) bool {
|
||||
_, err := os.Stat(path)
|
||||
return err == nil
|
||||
}
|
||||
|
||||
func DownloadIconCli() {
|
||||
cliCmd := flag.NewFlagSet("download", flag.ExitOnError)
|
||||
outputPath := cliCmd.String("output-path", "", "path to output folder")
|
||||
prompt := cliCmd.String("prompt", "", "recraft icon prompt")
|
||||
if err := cliCmd.Parse(os.Args[2:]); err != nil {
|
||||
log.Fatalf("Failed to parse flags: %v", err)
|
||||
}
|
||||
if _, err := DownloadIcon(*prompt, *outputPath); err != nil {
|
||||
log.Fatalln(err)
|
||||
}
|
||||
}
|
||||
|
||||
func DownloadIcon(prompt string, outputPath string) (string, error) {
|
||||
sanatizedPrompt := SanatizePrompt(prompt)
|
||||
svgFilePath := FindNoneCollisionFileName(sanatizedPrompt, "svg", outputPath)
|
||||
url, err := GenerateIcon(prompt)
|
||||
if err != nil {
|
||||
log.Printf("error generating icon: %v\n", err)
|
||||
return "", err
|
||||
}
|
||||
if err := DownloadFile(svgFilePath, url); err != nil {
|
||||
log.Printf("error downloading %s : %v\n", url, err)
|
||||
return "", err
|
||||
}
|
||||
log.Println("success")
|
||||
return svgFilePath, nil
|
||||
}
|
||||
|
||||
func FindNoneCollisionFileName(fileName, fileExt, path string) string {
|
||||
|
||||
filePath := filepath.Join(path, fileName+"."+fileExt)
|
||||
for fileCount := 1; FileExists(filePath); fileCount += 1 {
|
||||
log.Println("file exists: ", filePath)
|
||||
filePath = filepath.Join(
|
||||
path,
|
||||
fmt.Sprintf("%s (%d).%s", fileName, fileCount, fileExt),
|
||||
)
|
||||
}
|
||||
return filePath
|
||||
}
|
||||
|
||||
func CleanHeader(input string) string {
|
||||
re := regexp.MustCompile(`\s*(style|width|height)="[^"]*"`)
|
||||
return re.ReplaceAllString(input, "")
|
||||
}
|
||||
Reference in New Issue
Block a user