182 lines
3.5 KiB
Go
Raw Normal View History

2017-06-12 18:06:31 +00:00
package main
import (
2017-06-13 11:23:54 +00:00
"fmt"
2017-06-12 18:06:31 +00:00
"os"
"path"
2017-06-13 12:47:17 +00:00
"reflect"
2017-06-12 18:06:31 +00:00
"runtime"
2017-07-08 14:10:22 +00:00
"sync"
2017-06-12 18:06:31 +00:00
"time"
2017-06-13 12:47:17 +00:00
_ "image/gif"
_ "image/jpeg"
_ "image/png"
2017-06-15 14:56:06 +00:00
"github.com/aerogo/log"
2017-06-12 18:06:31 +00:00
"github.com/animenotifier/arn"
2017-07-18 01:55:47 +00:00
"github.com/animenotifier/avatar"
"github.com/animenotifier/avatar/outputs"
"github.com/animenotifier/avatar/sources"
2017-06-13 11:23:54 +00:00
"github.com/fatih/color"
2017-06-12 22:06:35 +00:00
)
const (
2017-06-13 15:06:30 +00:00
webPQuality = 80
2017-06-12 18:06:31 +00:00
)
2017-07-18 01:55:47 +00:00
var avatarSources []avatar.Source
var avatarOutputs []avatar.Output
2017-06-21 18:33:57 +00:00
var avatarLog = log.New()
2017-07-08 14:10:22 +00:00
var wg sync.WaitGroup
2017-06-13 11:23:54 +00:00
// Main
2017-06-12 18:06:31 +00:00
func main() {
2017-06-14 15:16:03 +00:00
color.Yellow("Generating user avatars")
2017-06-13 11:23:54 +00:00
// Switch to main directory
exe, err := os.Executable()
if err != nil {
panic(err)
}
root := path.Dir(exe)
os.Chdir(path.Join(root, "../../"))
2017-06-12 22:06:35 +00:00
2017-06-15 14:56:06 +00:00
// Log
avatarLog.AddOutput(log.File("logs/avatar.log"))
2017-06-15 15:03:22 +00:00
defer avatarLog.Flush()
2017-06-15 14:56:06 +00:00
2017-06-13 11:23:54 +00:00
// Define the avatar sources
2017-07-18 01:55:47 +00:00
avatarSources = []avatar.Source{
&sources.Gravatar{
2017-06-13 15:06:30 +00:00
Rating: "pg",
2017-07-18 01:55:47 +00:00
RequestLimiter: time.NewTicker(100 * time.Millisecond),
2017-06-13 15:06:30 +00:00
},
2017-07-18 01:55:47 +00:00
&sources.MyAnimeList{
2017-06-15 17:22:09 +00:00
RequestLimiter: time.NewTicker(250 * time.Millisecond),
2017-06-13 15:06:30 +00:00
},
2017-07-18 01:55:47 +00:00
&sources.FileSystem{
2017-07-08 15:31:04 +00:00
Directory: "images/avatars/large/",
},
2017-06-13 11:23:54 +00:00
}
2017-06-13 12:47:17 +00:00
// Define the avatar outputs
2017-07-18 01:55:47 +00:00
avatarOutputs = []avatar.Output{
2017-06-13 15:06:30 +00:00
// Original - Large
2017-07-18 01:55:47 +00:00
&outputs.OriginalFile{
2017-07-08 13:40:13 +00:00
Directory: "images/avatars/large/",
2017-06-13 12:47:17 +00:00
Size: arn.AvatarMaxSize,
},
2017-06-13 15:06:30 +00:00
// Original - Small
2017-07-18 01:55:47 +00:00
&outputs.OriginalFile{
2017-07-08 13:40:13 +00:00
Directory: "images/avatars/small/",
2017-06-13 12:47:17 +00:00
Size: arn.AvatarSmallSize,
},
2017-06-13 15:06:30 +00:00
// WebP - Large
2017-07-18 01:55:47 +00:00
&outputs.WebPFile{
2017-07-08 13:40:13 +00:00
Directory: "images/avatars/large/",
2017-06-13 12:47:17 +00:00
Size: arn.AvatarMaxSize,
Quality: webPQuality,
},
2017-06-13 15:06:30 +00:00
// WebP - Small
2017-07-18 01:55:47 +00:00
&outputs.WebPFile{
2017-07-08 13:40:13 +00:00
Directory: "images/avatars/small/",
2017-06-13 12:47:17 +00:00
Size: arn.AvatarSmallSize,
Quality: webPQuality,
},
}
if InvokeShellArgs() {
return
}
2017-06-13 11:23:54 +00:00
// Worker queue
2017-07-08 14:10:22 +00:00
usersQueue := make(chan *arn.User, runtime.NumCPU())
2017-06-13 15:06:30 +00:00
StartWorkers(usersQueue, Work)
2017-06-12 18:06:31 +00:00
2017-07-13 07:20:12 +00:00
allUsers, _ := arn.AllUsers()
2017-06-13 11:23:54 +00:00
// We'll send each user to one of the worker threads
2017-07-13 07:20:12 +00:00
for _, user := range allUsers {
2017-07-08 14:10:22 +00:00
wg.Add(1)
2017-06-13 11:23:54 +00:00
usersQueue <- user
}
2017-06-14 15:16:03 +00:00
2017-07-08 14:10:22 +00:00
wg.Wait()
2017-06-14 15:16:03 +00:00
color.Green("Finished.")
2017-06-13 11:23:54 +00:00
}
2017-06-12 22:06:35 +00:00
2017-06-13 11:23:54 +00:00
// StartWorkers creates multiple workers to handle a user each.
2017-06-13 15:06:30 +00:00
func StartWorkers(queue chan *arn.User, work func(*arn.User)) {
2017-06-13 11:23:54 +00:00
for w := 0; w < runtime.NumCPU(); w++ {
go func() {
for user := range queue {
work(user)
2017-07-08 14:19:31 +00:00
wg.Done()
2017-06-12 18:06:31 +00:00
}
2017-06-13 11:23:54 +00:00
}()
2017-06-12 18:06:31 +00:00
}
2017-06-13 11:23:54 +00:00
}
2017-06-12 18:06:31 +00:00
2017-06-13 11:23:54 +00:00
// Work handles a single user.
func Work(user *arn.User) {
2017-07-17 17:56:26 +00:00
user.Avatar.Extension = ""
2017-06-13 12:47:17 +00:00
2017-06-13 11:23:54 +00:00
for _, source := range avatarSources {
2017-07-18 01:55:47 +00:00
avatar, err := source.GetAvatar(user)
if err != nil {
avatarLog.Error(err)
continue
}
2017-06-13 11:23:54 +00:00
if avatar == nil {
2017-06-13 15:06:30 +00:00
// fmt.Println(color.RedString("✘"), reflect.TypeOf(source).Elem().Name(), user.Nick)
2017-06-13 11:23:54 +00:00
continue
}
2017-07-08 15:34:21 +00:00
// Name of source
2017-07-17 18:03:07 +00:00
user.Avatar.Source = reflect.TypeOf(source).Elem().Name()
2017-07-08 15:31:04 +00:00
2017-07-08 15:34:21 +00:00
// Log
2017-07-17 18:03:07 +00:00
fmt.Println(color.GreenString("✔"), user.Avatar.Source, "|", user.Nick, "|", avatar)
2017-07-08 15:34:21 +00:00
2017-07-17 18:03:07 +00:00
// Avoid JPG quality loss (if it's on the file system, we don't need to write it again)
if user.Avatar.Source == "FileSystem" {
2017-07-17 17:56:26 +00:00
user.Avatar.Extension = avatar.Extension()
2017-07-08 15:34:21 +00:00
break
2017-07-08 15:31:04 +00:00
}
2017-06-13 12:47:17 +00:00
for _, writer := range avatarOutputs {
err := writer.SaveAvatar(avatar)
2017-06-13 11:23:54 +00:00
2017-06-13 12:47:17 +00:00
if err != nil {
color.Red(err.Error())
}
2017-06-13 11:23:54 +00:00
}
2017-06-13 12:47:17 +00:00
break
2017-06-12 18:06:31 +00:00
}
2017-06-13 11:23:54 +00:00
2017-07-08 13:40:13 +00:00
// Since this a very long running job, refresh user data before saving it.
2017-07-17 17:56:26 +00:00
avatarExt := user.Avatar.Extension
2017-07-17 18:03:07 +00:00
avatarSrc := user.Avatar.Source
2017-07-08 13:40:13 +00:00
user, err := arn.GetUser(user.ID)
if err != nil {
avatarLog.Error("Can't refresh user info:", user.ID, user.Nick)
return
}
2017-06-13 12:47:17 +00:00
// Save avatar data
2017-07-17 17:56:26 +00:00
user.Avatar.Extension = avatarExt
2017-07-17 18:03:07 +00:00
user.Avatar.Source = avatarSrc
2017-06-13 12:47:17 +00:00
user.Save()
2017-06-12 18:06:31 +00:00
}