package scraper import ( "context" "fmt" "net/http" "strconv" "github.com/stashapp/stash/pkg/match" "github.com/stashapp/stash/pkg/models" "github.com/stashapp/stash/pkg/txn" ) // autoTagScraperID is the scraper ID for the built-in AutoTag scraper const ( autoTagScraperID = "builtin_autotag" autoTagScraperName = "Auto Tag" ) type autotagScraper struct { txnManager txn.Manager performerReader models.PerformerAutoTagQueryer studioReader models.StudioAutoTagQueryer tagReader models.TagAutoTagQueryer globalConfig GlobalConfig } func autotagMatchPerformers(ctx context.Context, path string, performerReader models.PerformerAutoTagQueryer, trimExt bool) ([]*models.ScrapedPerformer, error) { p, err := match.PathToPerformers(ctx, path, performerReader, nil, trimExt) 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, StoredID: &id, } if pp.Gender != nil && pp.Gender.IsValid() { v := pp.Gender.String() sp.Gender = &v } ret = append(ret, sp) } return ret, nil } func autotagMatchStudio(ctx context.Context, path string, studioReader models.StudioAutoTagQueryer, trimExt bool) (*models.ScrapedStudio, error) { studio, err := match.PathToStudio(ctx, path, studioReader, nil, trimExt) if err != nil { return nil, fmt.Errorf("error matching studios: %w", err) } if studio != nil { id := strconv.Itoa(studio.ID) return &models.ScrapedStudio{ Name: studio.Name, StoredID: &id, }, nil } return nil, nil } func autotagMatchTags(ctx context.Context, path string, tagReader models.TagAutoTagQueryer, trimExt bool) ([]*models.ScrapedTag, error) { t, err := match.PathToTags(ctx, path, tagReader, nil, trimExt) 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 } func (s autotagScraper) viaScene(ctx context.Context, _client *http.Client, scene *models.Scene) (*ScrapedScene, error) { var ret *ScrapedScene const trimExt = false // populate performers, studio and tags based on scene path if err := txn.WithReadTxn(ctx, s.txnManager, func(ctx context.Context) error { path := scene.Path if path == "" { return nil } performers, err := autotagMatchPerformers(ctx, path, s.performerReader, trimExt) if err != nil { return fmt.Errorf("autotag scraper viaScene: %w", err) } studio, err := autotagMatchStudio(ctx, path, s.studioReader, trimExt) if err != nil { return fmt.Errorf("autotag scraper viaScene: %w", err) } tags, err := autotagMatchTags(ctx, path, s.tagReader, trimExt) if err != nil { return fmt.Errorf("autotag scraper viaScene: %w", err) } if len(performers) > 0 || studio != nil || len(tags) > 0 { ret = &ScrapedScene{ Performers: performers, Studio: studio, Tags: tags, } } return nil }); err != nil { return nil, err } return ret, nil } func (s autotagScraper) viaGallery(ctx context.Context, _client *http.Client, gallery *models.Gallery) (*ScrapedGallery, error) { path := gallery.Path if path == "" { // not valid for non-path-based galleries return nil, nil } // only trim extension if gallery is file-based trimExt := gallery.PrimaryFileID != nil var ret *ScrapedGallery // populate performers, studio and tags based on scene path if err := txn.WithReadTxn(ctx, s.txnManager, func(ctx context.Context) error { path := gallery.Path performers, err := autotagMatchPerformers(ctx, path, s.performerReader, trimExt) if err != nil { return fmt.Errorf("autotag scraper viaGallery: %w", err) } studio, err := autotagMatchStudio(ctx, path, s.studioReader, trimExt) if err != nil { return fmt.Errorf("autotag scraper viaGallery: %w", err) } tags, err := autotagMatchTags(ctx, path, s.tagReader, trimExt) if err != nil { return fmt.Errorf("autotag scraper viaGallery: %w", err) } if len(performers) > 0 || studio != nil || len(tags) > 0 { ret = &ScrapedGallery{ Performers: performers, Studio: studio, Tags: tags, } } return nil }); err != nil { return nil, err } return ret, nil } func (s autotagScraper) supports(ty ScrapeContentType) bool { switch ty { case ScrapeContentTypeScene: return true case ScrapeContentTypeGallery: return true } return false } func (s autotagScraper) supportsURL(url string, ty ScrapeContentType) bool { return false } func (s autotagScraper) spec() Scraper { supportedScrapes := []ScrapeType{ ScrapeTypeFragment, } return Scraper{ ID: autoTagScraperID, Name: autoTagScraperName, Scene: &ScraperSpec{ SupportedScrapes: supportedScrapes, }, Gallery: &ScraperSpec{ SupportedScrapes: supportedScrapes, }, } } func getAutoTagScraper(repo Repository, globalConfig GlobalConfig) scraper { base := autotagScraper{ txnManager: repo.TxnManager, performerReader: repo.PerformerFinder, studioReader: repo.StudioFinder, tagReader: repo.TagFinder, globalConfig: globalConfig, } return base }