stash/pkg/scraper/autotag.go

package scraper

import (
	"context"
	"errors"
	"fmt"
	"strconv"

	"github.com/stashapp/stash/pkg/match"
	"github.com/stashapp/stash/pkg/models"
)

// autoTagScraperID is the scraper ID for the built-in AutoTag scraper
const (
	autoTagScraperID   = "builtin_autotag"
	autoTagScraperName = "Auto Tag"
)

var errNotSupported = errors.New("not supported")

type autotagScraper struct {
	txnManager   models.TransactionManager
	globalConfig GlobalConfig
}

func (s *autotagScraper) matchPerformers(path string, performerReader models.PerformerReader) ([]*models.ScrapedPerformer, error) {
	p, err := match.PathToPerformers(path, performerReader)
	if err != nil {
		return nil, fmt.Errorf("error matching performers: %w", err)
	}

	var ret []*models.ScrapedPerformer
	for _, pp := range p {
		id := strconv.Itoa(pp.ID)

		sp := &models.ScrapedPerformer{
			Name:     &pp.Name.String,
			StoredID: &id,
		}
		if pp.Gender.Valid {
			sp.Gender = &pp.Gender.String
		}

		ret = append(ret, sp)
	}

	return ret, nil
}

func (s *autotagScraper) matchStudio(path string, studioReader models.StudioReader) (*models.ScrapedStudio, error) {
	st, err := match.PathToStudios(path, studioReader)
	if err != nil {
		return nil, fmt.Errorf("error matching studios: %w", err)
	}

	if len(st) > 0 {
		id := strconv.Itoa(st[0].ID)
		return &models.ScrapedStudio{
			Name:     st[0].Name.String,
			StoredID: &id,
		}, nil
	}

	return nil, nil
}

func (s *autotagScraper) matchTags(path string, tagReader models.TagReader) ([]*models.ScrapedTag, error) {
	t, err := match.PathToTags(path, tagReader)
	if err != nil {
		return nil, fmt.Errorf("error matching tags: %w", err)
	}

	var ret []*models.ScrapedTag
	for _, tt := range t {
		id := strconv.Itoa(tt.ID)

		st := &models.ScrapedTag{
			Name:     tt.Name,
			StoredID: &id,
		}

		ret = append(ret, st)
	}

	return ret, nil
}

type autotagSceneScraper struct {
	*autotagScraper
}

func (c *autotagSceneScraper) scrapeByName(name string) ([]*models.ScrapedScene, error) {
	return nil, errNotSupported
}

func (c *autotagSceneScraper) scrapeByScene(scene *models.Scene) (*models.ScrapedScene, error) {
	var ret *models.ScrapedScene

	// populate performers, studio and tags based on scene path
	if err := c.txnManager.WithReadTxn(context.TODO(), func(r models.ReaderRepository) error {
		path := scene.Path
		performers, err := c.matchPerformers(path, r.Performer())
		if err != nil {
			return err
		}
		studio, err := c.matchStudio(path, r.Studio())
		if err != nil {
			return err
		}

		tags, err := c.matchTags(path, r.Tag())
		if err != nil {
			return err
		}

		if len(performers) > 0 || studio != nil || len(tags) > 0 {
			ret = &models.ScrapedScene{
				Performers: performers,
				Studio:     studio,
				Tags:       tags,
			}
		}

		return nil
	}); err != nil {
		return nil, err
	}

	return ret, nil
}

func (c *autotagSceneScraper) scrapeByFragment(scene models.ScrapedSceneInput) (*models.ScrapedScene, error) {
	return nil, errNotSupported
}

func (c *autotagSceneScraper) scrapeByURL(url string) (*models.ScrapedScene, error) {
	return nil, errNotSupported
}

type autotagGalleryScraper struct {
	*autotagScraper
}

func (c *autotagGalleryScraper) scrapeByGallery(gallery *models.Gallery) (*models.ScrapedGallery, error) {
	if !gallery.Path.Valid {
		// not valid for non-path-based galleries
		return nil, nil
	}

	var ret *models.ScrapedGallery

	// populate performers, studio and tags based on scene path
	if err := c.txnManager.WithReadTxn(context.TODO(), func(r models.ReaderRepository) error {
		path := gallery.Path.String
		performers, err := c.matchPerformers(path, r.Performer())
		if err != nil {
			return err
		}
		studio, err := c.matchStudio(path, r.Studio())
		if err != nil {
			return err
		}

		tags, err := c.matchTags(path, r.Tag())
		if err != nil {
			return err
		}

		if len(performers) > 0 || studio != nil || len(tags) > 0 {
			ret = &models.ScrapedGallery{
				Performers: performers,
				Studio:     studio,
				Tags:       tags,
			}
		}

		return nil
	}); err != nil {
		return nil, err
	}

	return ret, nil
}

func (c *autotagGalleryScraper) scrapeByFragment(gallery models.ScrapedGalleryInput) (*models.ScrapedGallery, error) {
	return nil, errNotSupported
}

func (c *autotagGalleryScraper) scrapeByURL(url string) (*models.ScrapedGallery, error) {
	return nil, errNotSupported
}

func getAutoTagScraper(txnManager models.TransactionManager, globalConfig GlobalConfig) scraper {
	base := autotagScraper{
		txnManager:   txnManager,
		globalConfig: globalConfig,
	}

	supportedScrapes := []models.ScrapeType{
		models.ScrapeTypeFragment,
	}

	return scraper{
		ID: autoTagScraperID,
		Spec: &models.Scraper{
			ID:   autoTagScraperID,
			Name: autoTagScraperName,
			Scene: &models.ScraperSpec{
				SupportedScrapes: supportedScrapes,
			},
			Gallery: &models.ScraperSpec{
				SupportedScrapes: supportedScrapes,
			},
		},
		Scene:   &autotagSceneScraper{&base},
		Gallery: &autotagGalleryScraper{&base},
	}
}