From 5da999ccd66f37c187c3b9f09dc7a44f9f9ea6aa Mon Sep 17 00:00:00 2001
From: Bart Mesuere <Bart.Mesuere@UGent.be>
Date: Tue, 2 Jul 2024 17:13:48 +0200
Subject: [PATCH] add a few utility methods

---
 lib/commands/unipept/pept2lca.ts           |  8 ++---
 lib/commands/unipept/unipept_subcommand.ts | 35 ++++++++++++++++++----
 2 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/lib/commands/unipept/pept2lca.ts b/lib/commands/unipept/pept2lca.ts
index 5fdc0239..97b8a871 100644
--- a/lib/commands/unipept/pept2lca.ts
+++ b/lib/commands/unipept/pept2lca.ts
@@ -24,11 +24,11 @@ The command will give priority to the first way tryptic peptides are passed, in
       .action((args, options) => this.run(args, options));
   }
 
-  defaultBatchSize(): number {
-    return 100;
+  requiredFields(): string[] {
+    return ["peptide"];
   }
 
-  async run(args: string[], options: object) {
-    await super.run(args, options);
+  defaultBatchSize(): number {
+    return 100;
   }
 }
diff --git a/lib/commands/unipept/unipept_subcommand.ts b/lib/commands/unipept/unipept_subcommand.ts
index 873a24bc..dca448ce 100644
--- a/lib/commands/unipept/unipept_subcommand.ts
+++ b/lib/commands/unipept/unipept_subcommand.ts
@@ -12,6 +12,8 @@ export abstract class UnipeptSubcommand {
   user_agent: string;
   host = "https://api.unipept.ugent.be";
   url?: string;
+  selectedFields?: RegExp[];
+  fasta: boolean;
 
   constructor(name: string) {
     this.name = name;
@@ -22,6 +24,10 @@ export abstract class UnipeptSubcommand {
   }
   abstract defaultBatchSize(): number;
 
+  requiredFields(): string[] {
+    return [];
+  }
+
   create(name: string): Command {
     const command = new Command(name);
 
@@ -47,7 +53,7 @@ export abstract class UnipeptSubcommand {
 
     for await (const input of this.getInputIterator(args, options.input)) {
       slice.push(input);
-      if (slice.length >= this.defaultBatchSize()) {
+      if (slice.length >= this.batchSize) {
         await this.processBatch(slice);
         slice = [];
       }
@@ -68,7 +74,7 @@ export abstract class UnipeptSubcommand {
   }
 
   private constructRequestBody(slice: string[]): URLSearchParams {
-    const names = this.constructSelectedFields().length === 0 || this.constructSelectedFields().some(regex => regex.toString().includes("name") || regex.toString().includes(".*$"));
+    const names = this.getSelectedFields().length === 0 || this.getSelectedFields().some(regex => regex.toString().includes("name") || regex.toString().includes(".*$"));
     return new URLSearchParams({
       input: JSON.stringify(slice),
       equate_il: this.options.equate,
@@ -77,9 +83,24 @@ export abstract class UnipeptSubcommand {
     });
   }
 
-  // TODO: implement
-  private constructSelectedFields(): RegExp[] {
-    return [];
+  private getSelectedFields(): RegExp[] {
+    if (this.selectedFields) return this.getSelectedFields();
+
+    const fields = (this.options.fields as string[]).flatMap(f => f.split(","));
+    if (this.fasta && fields.length > 0) {
+      fields.push(...this.requiredFields());
+    }
+    this.selectedFields = fields.map(f => this.globToRegex(f));
+
+    return this.selectedFields;
+  }
+
+  private get batchSize(): number {
+    if (this.options.batch) {
+      return +this.options.batch;
+    } else {
+      return this.defaultBatchSize();
+    }
   }
 
   /**
@@ -108,4 +129,8 @@ export abstract class UnipeptSubcommand {
       return `http://${host}`;
     }
   }
+
+  private globToRegex(glob: string): RegExp {
+    return new RegExp(glob.replace(/\*/g, ".*"));
+  }
 }