Skip to content

Commit

Permalink
Identify: Options to skip multiple results and single name performers (…
Browse files Browse the repository at this point in the history 
…stashapp#3707)

Co-authored-by: WithoutPants <53250216+WithoutPants@users.noreply.github.com>
  • Loading branch information
@Flashy78 @WithoutPants
Flashy78 and WithoutPants authored Jul 11, 2023
1 parent ff22577 commit cbdd4d3
Show file tree
Hide file tree
Showing 19 changed files with 581 additions and 136 deletions.
4 changes: 4 additions & 0 deletions graphql/documents/data/config.graphql
View file
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,10 @@ fragment IdentifyMetadataOptionsData on IdentifyMetadataOptions {
setCoverImage
setOrganized
includeMalePerformers
skipMultipleMatches
skipMultipleMatchTag
skipSingleNamePerformers
skipSingleNamePerformerTag
}

fragment ScraperSourceData on ScraperSource {
Expand Down
16 changes: 16 additions & 0 deletions graphql/schema/types/metadata.graphql
View file
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,14 @@ input IdentifyMetadataOptionsInput {
setOrganized: Boolean
"""defaults to true if not provided"""
includeMalePerformers: Boolean
"""defaults to true if not provided"""
skipMultipleMatches: Boolean
"""tag to tag skipped multiple matches with"""
skipMultipleMatchTag: String
"""defaults to true if not provided"""
skipSingleNamePerformers: Boolean
"""tag to tag skipped single name performers with"""
skipSingleNamePerformerTag: String
}

input IdentifySourceInput {
Expand Down Expand Up @@ -222,6 +230,14 @@ type IdentifyMetadataOptions {
setOrganized: Boolean
"""defaults to true if not provided"""
includeMalePerformers: Boolean
"""defaults to true if not provided"""
skipMultipleMatches: Boolean
"""tag to tag skipped multiple matches with"""
skipMultipleMatchTag: String
"""defaults to true if not provided"""
skipSingleNamePerformers: Boolean
"""tag to tag skipped single name performers with"""
skipSingleNamePerformerTag: String
}

type IdentifySource {
Expand Down
206 changes: 155 additions & 51 deletions internal/identify/identify.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,33 @@ package identify

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

"github.com/stashapp/stash/pkg/logger"
"github.com/stashapp/stash/pkg/models"
"github.com/stashapp/stash/pkg/scene"
"github.com/stashapp/stash/pkg/scraper"
"github.com/stashapp/stash/pkg/sliceutil/intslice"
"github.com/stashapp/stash/pkg/txn"
"github.com/stashapp/stash/pkg/utils"
)

var (
ErrSkipSingleNamePerformer = errors.New("a performer was skipped because they only had a single name and no disambiguation")
)

type MultipleMatchesFoundError struct {
Source ScraperSource
}

func (e *MultipleMatchesFoundError) Error() string {
return fmt.Sprintf("multiple matches found for %s", e.Source.Name)
}

type SceneScraper interface {
ScrapeScene(ctx context.Context, sceneID int) (*scraper.ScrapedScene, error)
ScrapeScenes(ctx context.Context, sceneID int) ([]*scraper.ScrapedScene, error)
}

type SceneUpdatePostHookExecutor interface {
Expand All @@ -31,21 +46,39 @@ type SceneIdentifier struct {
SceneReaderUpdater SceneReaderUpdater
StudioCreator StudioCreator
PerformerCreator PerformerCreator
TagCreator TagCreator
TagCreatorFinder TagCreatorFinder

DefaultOptions *MetadataOptions
Sources []ScraperSource
SceneUpdatePostHookExecutor SceneUpdatePostHookExecutor
}

func (t *SceneIdentifier) Identify(ctx context.Context, txnManager txn.Manager, scene *models.Scene) error {
result, err := t.scrapeScene(ctx, scene)
result, err := t.scrapeScene(ctx, txnManager, scene)
var multipleMatchErr *MultipleMatchesFoundError
if err != nil {
return err
if !errors.As(err, &multipleMatchErr) {
return err
}
}

if result == nil {
logger.Debugf("Unable to identify %s", scene.Path)
if multipleMatchErr != nil {
logger.Debugf("Identify skipped because multiple results returned for %s", scene.Path)

// find if the scene should be tagged for multiple results
options := t.getOptions(multipleMatchErr.Source)
if options.SkipMultipleMatchTag != nil && len(*options.SkipMultipleMatchTag) > 0 {
// Tag it with the multiple results tag
err := t.addTagToScene(ctx, txnManager, scene, *options.SkipMultipleMatchTag)
if err != nil {
return err
}
return nil
}
} else {
logger.Debugf("Unable to identify %s", scene.Path)
}
return nil
}

Expand All @@ -62,63 +95,98 @@ type scrapeResult struct {
source ScraperSource
}

func (t *SceneIdentifier) scrapeScene(ctx context.Context, scene *models.Scene) (*scrapeResult, error) {
func (t *SceneIdentifier) scrapeScene(ctx context.Context, txnManager txn.Manager, scene *models.Scene) (*scrapeResult, error) {
// iterate through the input sources
for _, source := range t.Sources {
// scrape using the source
scraped, err := source.Scraper.ScrapeScene(ctx, scene.ID)
results, err := source.Scraper.ScrapeScenes(ctx, scene.ID)
if err != nil {
logger.Errorf("error scraping from %v: %v", source.Scraper, err)
continue
}

// if results were found then return
if scraped != nil {
return &scrapeResult{
result: scraped,
source: source,
}, nil
if len(results) > 0 {
options := t.getOptions(source)
if len(results) > 1 && utils.IsTrue(options.SkipMultipleMatches) {
return nil, &MultipleMatchesFoundError{
Source: source,
}
} else {
// if results were found then return
return &scrapeResult{
result: results[0],
source: source,
}, nil
}
}
}

return nil, nil
}

// Returns a MetadataOptions object with any default options overwritten by source specific options
func (t *SceneIdentifier) getOptions(source ScraperSource) MetadataOptions {
options := *t.DefaultOptions
if source.Options == nil {
return options
}
if source.Options.SetCoverImage != nil {
options.SetCoverImage = source.Options.SetCoverImage
}
if source.Options.SetOrganized != nil {
options.SetOrganized = source.Options.SetOrganized
}
if source.Options.IncludeMalePerformers != nil {
options.IncludeMalePerformers = source.Options.IncludeMalePerformers
}
if source.Options.SkipMultipleMatches != nil {
options.SkipMultipleMatches = source.Options.SkipMultipleMatches
}
if source.Options.SkipMultipleMatchTag != nil && len(*source.Options.SkipMultipleMatchTag) > 0 {
options.SkipMultipleMatchTag = source.Options.SkipMultipleMatchTag
}
if source.Options.SkipSingleNamePerformers != nil {
options.SkipSingleNamePerformers = source.Options.SkipSingleNamePerformers
}
if source.Options.SkipSingleNamePerformerTag != nil && len(*source.Options.SkipSingleNamePerformerTag) > 0 {
options.SkipSingleNamePerformerTag = source.Options.SkipSingleNamePerformerTag
}
return options
}

func (t *SceneIdentifier) getSceneUpdater(ctx context.Context, s *models.Scene, result *scrapeResult) (*scene.UpdateSet, error) {
ret := &scene.UpdateSet{
ID: s.ID,
}

options := []MetadataOptions{}
allOptions := []MetadataOptions{}
if result.source.Options != nil {
options = append(options, *result.source.Options)
allOptions = append(allOptions, *result.source.Options)
}
if t.DefaultOptions != nil {
options = append(options, *t.DefaultOptions)
allOptions = append(allOptions, *t.DefaultOptions)
}

fieldOptions := getFieldOptions(options)

setOrganized := false
for _, o := range options {
if o.SetOrganized != nil {
setOrganized = *o.SetOrganized
break
}
}
fieldOptions := getFieldOptions(allOptions)
options := t.getOptions(result.source)

scraped := result.result

rel := sceneRelationships{
sceneReader: t.SceneReaderUpdater,
studioCreator: t.StudioCreator,
performerCreator: t.PerformerCreator,
tagCreator: t.TagCreator,
scene: s,
result: result,
fieldOptions: fieldOptions,
sceneReader: t.SceneReaderUpdater,
studioCreator: t.StudioCreator,
performerCreator: t.PerformerCreator,
tagCreatorFinder: t.TagCreatorFinder,
scene: s,
result: result,
fieldOptions: fieldOptions,
skipSingleNamePerformers: *options.SkipSingleNamePerformers,
}

setOrganized := false
if options.SetOrganized != nil {
setOrganized = *options.SetOrganized
}
ret.Partial = getScenePartial(s, scraped, fieldOptions, setOrganized)

studioID, err := rel.studio(ctx)
Expand All @@ -130,17 +198,19 @@ func (t *SceneIdentifier) getSceneUpdater(ctx context.Context, s *models.Scene,
ret.Partial.StudioID = models.NewOptionalInt(*studioID)
}

ignoreMale := false
for _, o := range options {
if o.IncludeMalePerformers != nil {
ignoreMale = !*o.IncludeMalePerformers
break
}
includeMalePerformers := true
if options.IncludeMalePerformers != nil {
includeMalePerformers = *options.IncludeMalePerformers
}

performerIDs, err := rel.performers(ctx, ignoreMale)
addSkipSingleNamePerformerTag := false
performerIDs, err := rel.performers(ctx, !includeMalePerformers)
if err != nil {
return nil, err
if errors.Is(err, ErrSkipSingleNamePerformer) {
addSkipSingleNamePerformerTag = true
} else {
return nil, err
}
}
if performerIDs != nil {
ret.Partial.PerformerIDs = &models.UpdateIDs{
Expand All @@ -153,6 +223,14 @@ func (t *SceneIdentifier) getSceneUpdater(ctx context.Context, s *models.Scene,
if err != nil {
return nil, err
}
if addSkipSingleNamePerformerTag {
tagID, err := strconv.ParseInt(*options.SkipSingleNamePerformerTag, 10, 64)
if err != nil {
return nil, fmt.Errorf("error converting tag ID %s: %w", *options.SkipSingleNamePerformerTag, err)
}

tagIDs = intslice.IntAppendUnique(tagIDs, int(tagID))
}
if tagIDs != nil {
ret.Partial.TagIDs = &models.UpdateIDs{
IDs: tagIDs,
Expand All @@ -171,15 +249,7 @@ func (t *SceneIdentifier) getSceneUpdater(ctx context.Context, s *models.Scene,
}
}

setCoverImage := false
for _, o := range options {
if o.SetCoverImage != nil {
setCoverImage = *o.SetCoverImage
break
}
}

if setCoverImage {
if options.SetCoverImage != nil && *options.SetCoverImage {
ret.CoverImage, err = rel.cover(ctx)
if err != nil {
return nil, err
Expand Down Expand Up @@ -241,6 +311,41 @@ func (t *SceneIdentifier) modifyScene(ctx context.Context, txnManager txn.Manage
return nil
}

func (t *SceneIdentifier) addTagToScene(ctx context.Context, txnManager txn.Manager, s *models.Scene, tagToAdd string) error {
if err := txn.WithTxn(ctx, txnManager, func(ctx context.Context) error {
tagID, err := strconv.Atoi(tagToAdd)
if err != nil {
return fmt.Errorf("error converting tag ID %s: %w", tagToAdd, err)
}

if err := s.LoadTagIDs(ctx, t.SceneReaderUpdater); err != nil {
return err
}
existing := s.TagIDs.List()

if intslice.IntInclude(existing, tagID) {
// skip if the scene was already tagged
return nil
}

if err := scene.AddTag(ctx, t.SceneReaderUpdater, s, tagID); err != nil {
return err
}

ret, err := t.TagCreatorFinder.Find(ctx, tagID)
if err != nil {
logger.Infof("Added tag id %s to skipped scene %s", tagToAdd, s.Path)
} else {
logger.Infof("Added tag %s to skipped scene %s", ret.Name, s.Path)
}

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

func getFieldOptions(options []MetadataOptions) map[string]*FieldOptions {
// prefer source-specific field strategies, then the defaults
ret := make(map[string]*FieldOptions)
Expand Down Expand Up @@ -291,8 +396,7 @@ func getScenePartial(scene *models.Scene, scraped *scraper.ScrapedScene, fieldOp
}

if setOrganized && !scene.Organized {
// just reuse the boolean since we know it's true
partial.Organized = models.NewOptionalBool(setOrganized)
partial.Organized = models.NewOptionalBool(true)
}

return partial
Expand Down
Loading

0 comments on commit cbdd4d3

Please sign in to comment.