Autotag scraper (#1817)

* Refactor scraper structures
* Move matching code into new package
* Add autotag scraper
* Always check first letter of auto-tag names
* Account for nulls

Co-authored-by: Kermie <kermie@isinthe.house>
This commit is contained in:
WithoutPants
2021-10-11 23:06:06 +11:00
committed by GitHub
parent b5381ff071
commit e9d48683f8
22 changed files with 1023 additions and 660 deletions

View File

@@ -9,8 +9,6 @@ import (
"strings"
"gopkg.in/yaml.v2"
"github.com/stashapp/stash/pkg/models"
)
type config struct {
@@ -194,7 +192,7 @@ type scraperDriverOptions struct {
Headers []*header `yaml:"headers"`
}
func loadScraperFromYAML(id string, reader io.Reader) (*config, error) {
func loadConfigFromYAML(id string, reader io.Reader) (*config, error) {
ret := &config{}
parser := yaml.NewDecoder(reader)
@@ -213,7 +211,7 @@ func loadScraperFromYAML(id string, reader io.Reader) (*config, error) {
return ret, nil
}
func loadScraperFromYAMLFile(path string) (*config, error) {
func loadConfigFromYAMLFile(path string) (*config, error) {
file, err := os.Open(path)
if err != nil {
return nil, err
@@ -224,7 +222,7 @@ func loadScraperFromYAMLFile(path string) (*config, error) {
id := filepath.Base(path)
id = id[:strings.LastIndex(id, ".")]
ret, err := loadScraperFromYAML(id, file)
ret, err := loadConfigFromYAML(id, file)
if err != nil {
return nil, err
}
@@ -234,78 +232,6 @@ func loadScraperFromYAMLFile(path string) (*config, error) {
return ret, nil
}
func (c config) toScraper() *models.Scraper {
ret := models.Scraper{
ID: c.ID,
Name: c.Name,
}
performer := models.ScraperSpec{}
if c.PerformerByName != nil {
performer.SupportedScrapes = append(performer.SupportedScrapes, models.ScrapeTypeName)
}
if c.PerformerByFragment != nil {
performer.SupportedScrapes = append(performer.SupportedScrapes, models.ScrapeTypeFragment)
}
if len(c.PerformerByURL) > 0 {
performer.SupportedScrapes = append(performer.SupportedScrapes, models.ScrapeTypeURL)
for _, v := range c.PerformerByURL {
performer.Urls = append(performer.Urls, v.URL...)
}
}
if len(performer.SupportedScrapes) > 0 {
ret.Performer = &performer
}
scene := models.ScraperSpec{}
if c.SceneByFragment != nil {
scene.SupportedScrapes = append(scene.SupportedScrapes, models.ScrapeTypeFragment)
}
if c.SceneByName != nil && c.SceneByQueryFragment != nil {
scene.SupportedScrapes = append(scene.SupportedScrapes, models.ScrapeTypeName)
}
if len(c.SceneByURL) > 0 {
scene.SupportedScrapes = append(scene.SupportedScrapes, models.ScrapeTypeURL)
for _, v := range c.SceneByURL {
scene.Urls = append(scene.Urls, v.URL...)
}
}
if len(scene.SupportedScrapes) > 0 {
ret.Scene = &scene
}
gallery := models.ScraperSpec{}
if c.GalleryByFragment != nil {
gallery.SupportedScrapes = append(gallery.SupportedScrapes, models.ScrapeTypeFragment)
}
if len(c.GalleryByURL) > 0 {
gallery.SupportedScrapes = append(gallery.SupportedScrapes, models.ScrapeTypeURL)
for _, v := range c.GalleryByURL {
gallery.Urls = append(gallery.Urls, v.URL...)
}
}
if len(gallery.SupportedScrapes) > 0 {
ret.Gallery = &gallery
}
movie := models.ScraperSpec{}
if len(c.MovieByURL) > 0 {
movie.SupportedScrapes = append(movie.SupportedScrapes, models.ScrapeTypeURL)
for _, v := range c.MovieByURL {
movie.Urls = append(movie.Urls, v.URL...)
}
}
if len(movie.SupportedScrapes) > 0 {
ret.Movie = &movie
}
return &ret
}
func (c config) supportsPerformers() bool {
return c.PerformerByName != nil || c.PerformerByFragment != nil || len(c.PerformerByURL) > 0
}
@@ -320,47 +246,6 @@ func (c config) matchesPerformerURL(url string) bool {
return false
}
func (c config) ScrapePerformerNames(name string, txnManager models.TransactionManager, globalConfig GlobalConfig) ([]*models.ScrapedPerformer, error) {
if c.PerformerByName != nil {
s := getScraper(*c.PerformerByName, txnManager, c, globalConfig)
return s.scrapePerformersByName(name)
}
return nil, nil
}
func (c config) ScrapePerformer(scrapedPerformer models.ScrapedPerformerInput, txnManager models.TransactionManager, globalConfig GlobalConfig) (*models.ScrapedPerformer, error) {
if c.PerformerByFragment != nil {
s := getScraper(*c.PerformerByFragment, txnManager, c, globalConfig)
return s.scrapePerformerByFragment(scrapedPerformer)
}
// try to match against URL if present
if scrapedPerformer.URL != nil && *scrapedPerformer.URL != "" {
return c.ScrapePerformerURL(*scrapedPerformer.URL, txnManager, globalConfig)
}
return nil, nil
}
func (c config) ScrapePerformerURL(url string, txnManager models.TransactionManager, globalConfig GlobalConfig) (*models.ScrapedPerformer, error) {
for _, scraper := range c.PerformerByURL {
if scraper.matchesURL(url) {
s := getScraper(scraper.scraperTypeConfig, txnManager, c, globalConfig)
ret, err := s.scrapePerformerByURL(url)
if err != nil {
return nil, err
}
if ret != nil {
return ret, nil
}
}
}
return nil, nil
}
func (c config) supportsScenes() bool {
return (c.SceneByName != nil && c.SceneByQueryFragment != nil) || c.SceneByFragment != nil || len(c.SceneByURL) > 0
}
@@ -401,103 +286,3 @@ func (c config) matchesMovieURL(url string) bool {
return false
}
func (c config) ScrapeSceneQuery(name string, txnManager models.TransactionManager, globalConfig GlobalConfig) ([]*models.ScrapedScene, error) {
if c.SceneByName != nil {
s := getScraper(*c.SceneByName, txnManager, c, globalConfig)
return s.scrapeScenesByName(name)
}
return nil, nil
}
func (c config) ScrapeSceneByScene(scene *models.Scene, txnManager models.TransactionManager, globalConfig GlobalConfig) (*models.ScrapedScene, error) {
if c.SceneByFragment != nil {
s := getScraper(*c.SceneByFragment, txnManager, c, globalConfig)
return s.scrapeSceneByScene(scene)
}
return nil, nil
}
func (c config) ScrapeSceneByFragment(scene models.ScrapedSceneInput, txnManager models.TransactionManager, globalConfig GlobalConfig) (*models.ScrapedScene, error) {
if c.SceneByQueryFragment != nil {
s := getScraper(*c.SceneByQueryFragment, txnManager, c, globalConfig)
return s.scrapeSceneByFragment(scene)
}
return nil, nil
}
func (c config) ScrapeSceneURL(url string, txnManager models.TransactionManager, globalConfig GlobalConfig) (*models.ScrapedScene, error) {
for _, scraper := range c.SceneByURL {
if scraper.matchesURL(url) {
s := getScraper(scraper.scraperTypeConfig, txnManager, c, globalConfig)
ret, err := s.scrapeSceneByURL(url)
if err != nil {
return nil, err
}
if ret != nil {
return ret, nil
}
}
}
return nil, nil
}
func (c config) ScrapeGalleryByGallery(gallery *models.Gallery, txnManager models.TransactionManager, globalConfig GlobalConfig) (*models.ScrapedGallery, error) {
if c.GalleryByFragment != nil {
s := getScraper(*c.GalleryByFragment, txnManager, c, globalConfig)
return s.scrapeGalleryByGallery(gallery)
}
return nil, nil
}
func (c config) ScrapeGalleryByFragment(gallery models.ScrapedGalleryInput, txnManager models.TransactionManager, globalConfig GlobalConfig) (*models.ScrapedGallery, error) {
if c.GalleryByFragment != nil {
// TODO - this should be galleryByQueryFragment
s := getScraper(*c.GalleryByFragment, txnManager, c, globalConfig)
return s.scrapeGalleryByFragment(gallery)
}
return nil, nil
}
func (c config) ScrapeGalleryURL(url string, txnManager models.TransactionManager, globalConfig GlobalConfig) (*models.ScrapedGallery, error) {
for _, scraper := range c.GalleryByURL {
if scraper.matchesURL(url) {
s := getScraper(scraper.scraperTypeConfig, txnManager, c, globalConfig)
ret, err := s.scrapeGalleryByURL(url)
if err != nil {
return nil, err
}
if ret != nil {
return ret, nil
}
}
}
return nil, nil
}
func (c config) ScrapeMovieURL(url string, txnManager models.TransactionManager, globalConfig GlobalConfig) (*models.ScrapedMovie, error) {
for _, scraper := range c.MovieByURL {
if scraper.matchesURL(url) {
s := getScraper(scraper.scraperTypeConfig, txnManager, c, globalConfig)
ret, err := s.scrapeMovieByURL(url)
if err != nil {
return nil, err
}
if ret != nil {
return ret, nil
}
}
}
return nil, nil
}