Refactor taxonomy handling and improve performance

This update introduces logic for handling embedded taxonomy, allowing users to specify taxonomy or extract it from sequences, while optimizing performance through reduced memory usage. It also adds batch reordering capabilities and implements comprehensive testing to ensure correct precedence rules and robust error handling for missing taxonomic paths.

pkg/obiiter/batchiterator.go

@@ -465,7 +465,9 @@ func (iterator IBioSequence) RebatchBySize(maxBytes int, maxCount int) IBioSeque
 
 	go func() {
 		order := 0
-		iterator = iterator.SortBatches()
+		// a local variable: reassigning the captured iterator would race with
+		// the IsPaired call below, made by the caller on the same variable
+		sorted := iterator.SortBatches()
 		buffer := obiseq.MakeBioSequenceSlice()
 		bufBytes := 0
 		source := ""
@@ -479,8 +481,8 @@ func (iterator IBioSequence) RebatchBySize(maxBytes int, maxCount int) IBioSeque
 			}
 		}
 
-		for iterator.Next() {
-			seqs := iterator.Get()
+		for sorted.Next() {
+			seqs := sorted.Get()
 			source = seqs.Source()
 			for _, s := range seqs.Slice() {
 				sz := s.MemorySize()

pkg/obitools/obiprimers/embedded_test.go

@@ -0,0 +1,141 @@
+package obiprimers
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiiter"
+	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiprimer"
+	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiseq"
+	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obitax"
+)
+
+// Three sequences below two families (5 and 6) of the same root, annotated like
+// obiannotate --taxonomic-path does.
+const embeddedFasta = `>s1 {"taxid":"taxon:10 [Alpha beta]@species","taxonomic_path":["taxon:1 [root]@no rank","taxon:5 [Famone]@family","taxon:10 [Alpha beta]@species"]}
+acgtacgtacgtacgtacgtacgtacgtacgtacgtacgt
+>s2 {"taxid":"taxon:11 [Alpha gamma]@species","taxonomic_path":["taxon:1 [root]@no rank","taxon:5 [Famone]@family","taxon:11 [Alpha gamma]@species"]}
+acgtacgtacgtacgtacgtacgtacgtacgtacgtacgt
+>s3 {"taxid":"taxon:20 [Delta eps]@species","taxonomic_path":["taxon:1 [root]@no rank","taxon:6 [Famtwo]@family","taxon:20 [Delta eps]@species"]}
+acgtacgtacgtacgtacgtacgtacgtacgtacgtacgt
+`
+
+func writeTemp(t *testing.T, content string) string {
+	t.Helper()
+	path := filepath.Join(t.TempDir(), "seqs.fasta")
+	if err := os.WriteFile(path, []byte(content), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	return path
+}
+
+func drain(it obiiter.IBioSequence) (n int) {
+	for it.Next() {
+		batch := it.Get()
+		n += len(batch.Slice())
+	}
+	return
+}
+
+func TestEmbeddedTaxonomyFromFiles(t *testing.T) {
+	path := writeTemp(t, embeddedFasta)
+
+	sequences, taxonomy, err := readSequencesAndTaxonomy([]string{path}, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if taxonomy == nil || taxonomy.Len() != 6 { // root, 2 families, 3 species
+		t.Fatalf("expected an embedded taxonomy of 6 taxa, got %v", taxonomy)
+	}
+
+	// the sequences are read in full on the second pass
+	restrict, err := obiprimer.BelongingTo(taxonomy, []string{"5"})
+	if err != nil {
+		t.Fatalf("a taxid of the embedded taxonomy must be usable with -r: %v", err)
+	}
+	corpus, err := obiprimer.LoadCorpus(sequences, obiprimer.Selection{Restrict: restrict})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if c := corpus.Counts(); c.Examples != 2 || c.Ignored != 1 {
+		t.Errorf("restricting to family 5: %+v", c)
+	}
+
+	taxa, err := corpus.AssignTaxa(taxonomy, "species")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if taxa.ExampleTaxa != 2 {
+		t.Errorf("%d example species, want 2", taxa.ExampleTaxa)
+	}
+}
+
+func TestEmbeddedTaxonomyFromMemory(t *testing.T) { // the standard input case
+	var slice obiseq.BioSequenceSlice
+	for _, s := range []struct{ id, taxid string }{{"a", "10"}, {"b", "20"}} {
+		seq := obiseq.NewBioSequence(s.id, []byte("acgtacgtacgt"), "")
+		seq.SetAttribute("taxonomic_path", []string{
+			"taxon:1 [root]@no rank", "taxon:" + s.taxid + " [Sp" + s.taxid + "]@species"})
+		slice = append(slice, seq)
+	}
+
+	sequences, taxonomy, err := embeddedTaxonomyInMemory(obiiter.IBatchOver("stdin", slice, 10))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if taxonomy.Len() != 3 {
+		t.Errorf("%d taxa, want 3", taxonomy.Len())
+	}
+	if n := drain(sequences); n != 2 {
+		t.Errorf("the sequences must be given back intact: %d of 2", n)
+	}
+}
+
+func TestMissingEmbeddedTaxonomyGivesAdvice(t *testing.T) {
+	path := writeTemp(t, ">s1 {\"taxid\":\"taxon:10 [Alpha beta]@species\"}\nacgtacgtacgtacgt\n")
+	_, _, err := readSequencesAndTaxonomy([]string{path}, nil)
+	if err == nil {
+		t.Fatal("sequences without a taxonomic path and no -t must be refused")
+	}
+	for _, want := range []string{"-t", "obiannotate --taxonomic-path"} {
+		if !strings.Contains(err.Error(), want) {
+			t.Errorf("the error should mention %q: %v", want, err)
+		}
+	}
+
+	var slice obiseq.BioSequenceSlice
+	slice = append(slice, obiseq.NewBioSequence("a", []byte("acgt"), ""))
+	if _, _, err := embeddedTaxonomyInMemory(obiiter.IBatchOver("stdin", slice, 10)); err == nil ||
+		!strings.Contains(err.Error(), "obiannotate --taxonomic-path") {
+		t.Errorf("the in-memory case must give the same advice: %v", err)
+	}
+}
+
+func TestGivenTaxonomyHasPrecedence(t *testing.T) {
+	path := writeTemp(t, embeddedFasta)
+	_, embedded, err := readSequencesAndTaxonomy([]string{path}, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// With a taxonomy given (-t), the embedded paths are ignored, even if present.
+	other := obitax.NewTaxonomy("other", "other", obitax.DefaultTaxidAlphabet)
+	sequences, used, err := readSequencesAndTaxonomy([]string{path}, other)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if used != other || used == embedded {
+		t.Error("the given taxonomy must be returned as is")
+	}
+	if n := drain(sequences); n != 3 {
+		t.Errorf("%d sequences read, want 3", n)
+	}
+}
+
+func TestPrimerPairsNeedsATaxonomy(t *testing.T) {
+	if _, err := CLIPrimerPairs(obiiter.NilIBioSequence, nil); err == nil {
+		t.Error("a nil taxonomy must be refused")
+	}
+}

pkg/obitools/obiprimers/obiprimers.go

@@ -3,19 +3,92 @@ package obiprimers
 import (
 	"fmt"
 
+	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obidefault"
 	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiiter"
 	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiprimer"
 	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obitax"
+	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obitools/obiconvert"
 	log "github.com/sirupsen/logrus"
 )
 
+const noTaxonomyAdvice = "no taxonomy: use the -t option, or give sequences annotated with their " +
+	"taxonomic path (obiannotate --taxonomic-path)"
+
+// embeddedTaxonomyInMemory extracts the taxonomy from the taxonomic paths
+// carried by the sequences, which it has to hold in memory since the iterator
+// cannot be read twice. It returns an iterator over the same sequences.
+func embeddedTaxonomyInMemory(sequences obiiter.IBioSequence) (obiiter.IBioSequence, *obitax.Taxonomy, error) {
+	source, slice := sequences.Load()
+	taxonomy, err := slice.ExtractTaxonomy(nil, false)
+	if err != nil {
+		return obiiter.NilIBioSequence, nil, fmt.Errorf("%s: %w", noTaxonomyAdvice, err)
+	}
+	return obiiter.IBatchOver(source, slice, obidefault.BatchSize()), taxonomy, nil
+}
+
+// readSequencesAndTaxonomy returns the sequences and the taxonomy to use: the
+// given one when there is one (the -t option) and otherwise the taxonomy
+// embedded in the sequences. Like obitag, it needs no -t option when the
+// sequence files were annotated with obiannotate --taxonomic-path.
+//
+// For files the embedded taxonomy is collected by a first pass over them that
+// keeps nothing but the taxonomic paths, and the sequences are then read a
+// second time: the selection of the examples needs the taxonomy, so reading in
+// one pass would mean holding every sequence in memory. Standard input cannot
+// be read twice, its sequences are held in memory.
+func readSequencesAndTaxonomy(args []string, given *obitax.Taxonomy) (obiiter.IBioSequence, *obitax.Taxonomy, error) {
+	if given != nil {
+		sequences, err := obiconvert.CLIReadBioSequences(args...)
+		return sequences, given, err
+	}
+
+	if len(args) == 0 {
+		log.Infof("No taxonomy given: reading the standard input in memory to extract its embedded taxonomy")
+		sequences, err := obiconvert.CLIReadBioSequences()
+		if err != nil {
+			return obiiter.NilIBioSequence, nil, err
+		}
+		return embeddedTaxonomyInMemory(sequences)
+	}
+
+	log.Infof("No taxonomy given: extracting the taxonomy embedded in the sequence files (first pass)")
+	first, err := obiconvert.CLIReadBioSequences(args...)
+	if err != nil {
+		return obiiter.NilIBioSequence, nil, err
+	}
+	taxonomy, err := first.ExtractTaxonomy(false)
+	if err != nil || taxonomy == nil {
+		if err == nil {
+			err = fmt.Errorf("the sequences carry no taxonomic path")
+		}
+		return obiiter.NilIBioSequence, nil, fmt.Errorf("%s: %w", noTaxonomyAdvice, err)
+	}
+	log.Infof("Embedded taxonomy extracted: %d taxa", taxonomy.Len())
+
+	sequences, err := obiconvert.CLIReadBioSequences(args...)
+	return sequences, taxonomy, err
+}
+
+// CLIReadSequencesAndTaxonomy reads the sequences of the command line and
+// returns them with the taxonomy to use, set as the default taxonomy: the one
+// of the -t option, or else the one embedded in the sequences.
+func CLIReadSequencesAndTaxonomy(args []string) (obiiter.IBioSequence, *obitax.Taxonomy, error) {
+	given := obitax.DefaultTaxonomy()
+	sequences, taxonomy, err := readSequencesAndTaxonomy(args, given)
+	if err != nil {
+		return obiiter.NilIBioSequence, nil, err
+	}
+	if given == nil {
+		taxonomy.SetAsDefault()
+	}
+	return sequences, taxonomy, nil
+}
+
 // CLIPrimerPairs reads the sequences, selects the examples and the
 // counterexamples from the taxonomy options, and searches the primer pairs.
-// The default taxonomy must be loaded.
-func CLIPrimerPairs(sequences obiiter.IBioSequence) ([]obiprimer.PairStats, error) {
-	taxonomy := obitax.DefaultTaxonomy()
+func CLIPrimerPairs(sequences obiiter.IBioSequence, taxonomy *obitax.Taxonomy) ([]obiprimer.PairStats, error) {
 	if taxonomy == nil {
-		return nil, fmt.Errorf("a taxonomy is required: use the -t option")
+		return nil, fmt.Errorf(noTaxonomyAdvice)
 	}
 
 	sel := obiprimer.Selection{}

pkg/obitools/obiprimers/options.go

@@ -46,7 +46,7 @@ var _partitions = 256
 func OptionSet(options *getoptions.GetOpt) {
 	obiconvert.InputOptionSet(options)
 	obiconvert.OutputModeOptionSet(options, true)
-	obioptions.LoadTaxonomyOptionSet(options, true, false)
+	obioptions.LoadTaxonomyOptionSet(options, false, false)
 
 	options.StringSliceVar(&_restrict, "restrict-to-taxon", 1, 1,
 		options.Alias("r"),