yt-emoji/main.go

124 lines
2.3 KiB
Go

package main
import (
"context"
"encoding/json"
"fmt"
"io"
"log"
"net/http"
"os"
"path/filepath"
"regexp"
"strings"
"github.com/alexflint/go-arg"
"github.com/schollz/progressbar/v3"
"golang.org/x/sync/errgroup"
)
var reSize = regexp.MustCompile(`\b([whs])\d+\b`)
var args struct {
Folder string `arg:"positional,required"`
Size int `default:"-1"`
Threads int `default:"-1"`
}
type Item struct {
Source string `json:"src"`
Label string `json:"alt"`
}
func main() {
log.SetFlags(0)
arg.MustParse(&args)
log.Print("Opening data file...")
dataPath := filepath.Join(args.Folder, "data.json")
dataFile, err := os.Open(dataPath)
if err != nil {
log.Fatal(err)
}
defer dataFile.Close()
log.Print("Decoding data...")
var data []Item
dec := json.NewDecoder(dataFile)
dec.DisallowUnknownFields()
if err := dec.Decode(&data); err != nil {
log.Fatal(err)
}
log.Printf("Found %d item(s)", len(data))
log.Print("Converting badges...")
badgeCount := 0
for i := range data {
if strings.ToLower(data[i].Label) == "custom badge for members" {
badgeCount++
data[i].Label = fmt.Sprintf("badge%d", badgeCount)
}
}
log.Printf("Found %d badge(s)", badgeCount)
group, _ := errgroup.WithContext(context.Background())
group.SetLimit(args.Threads)
bar := progressbar.Default(int64(len(data)), "Downloading...")
for _, todo := range data {
group.Go(func() error {
defer bar.Add(1)
var trueSource string
if args.Size >= 0 {
replacement := fmt.Sprintf("${1}%d", args.Size)
trueSource = reSize.ReplaceAllString(todo.Source, replacement)
} else {
parts := strings.SplitN(todo.Source, "=", 2)
if len(parts) != 2 {
return fmt.Errorf("source doesn't contain '='")
}
trueSource = parts[0]
}
resp, err := http.Get(trueSource)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != 200 {
return fmt.Errorf("unexpected status: %s", resp.Status)
}
name := fmt.Sprintf("%s_%s.png", args.Folder, todo.Label)
path := filepath.Join(args.Folder, name)
f, err := os.Create(path)
if err != nil {
return err
}
defer f.Close()
_, err = io.Copy(f, resp.Body)
return err
})
}
if err := group.Wait(); err != nil {
bar.Exit()
log.Fatal(err)
}
bar.Finish()
log.Print("Done")
}