Update app.py

app.py

@@ -7,7 +7,9 @@ import zipfile
 import gradio as gr
 
 # Ensure repo root is importable on Spaces
-sys.path.append(os.path.dirname(__file__))
+ROOT = os.path.dirname(__file__)
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
 
 import kmer_predict  # must be in repo root
 
@@ -17,6 +19,7 @@ FASTA_EXTS = (".fa", ".fasta", ".fas", ".fna")
 
 
 def _zip_dir(folder: str, zip_path: str) -> None:
+    """Zip the contents of folder into zip_path."""
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
         for root, _, files in os.walk(folder):
             for fn in files:
@@ -24,27 +27,67 @@ def _zip_dir(folder: str, zip_path: str) -> None:
                 rel = os.path.relpath(full, folder)
                 z.write(full, rel)
 
+
 def _safe_extract_zip(zip_path: str, dst_dir: str) -> None:
-    """Safely extract FASTA files from ZIP (prevents zip-slip)."""
+    """Safely extract only FASTA files from ZIP (prevents zip-slip)."""
     with zipfile.ZipFile(zip_path, "r") as z:
         for member in z.infolist():
             if member.is_dir():
                 continue
 
+            # Zip-slip protection
             target = os.path.normpath(os.path.join(dst_dir, member.filename))
             if not target.startswith(os.path.abspath(dst_dir) + os.sep):
                 continue
 
-            if not member.filename.lower().endswith((".fa", ".fasta", ".fas", ".fna")):
+            # Only FASTA-like files
+            if not member.filename.lower().endswith(FASTA_EXTS):
                 continue
 
             os.makedirs(os.path.dirname(target), exist_ok=True)
             with z.open(member) as src, open(target, "wb") as out:
                 shutil.copyfileobj(src, out)
 
-def run_prediction(unknown_files, kmer_zip, seqtype, mode, identity, coverage, fdr):
-    if not unknown_files:
-        raise gr.Error("Please upload at least one unknown FASTA file.")
+
+def _ingest_unknown_uploads(unknown_uploads, unknown_dir: str) -> None:
+    """
+    Accept unknown sequences as:
+      - FASTA files, and/or
+      - ZIP files containing FASTA files.
+    Copies/extracts into unknown_dir.
+    """
+    os.makedirs(unknown_dir, exist_ok=True)
+
+    if not unknown_uploads:
+        return
+
+    for idx, f in enumerate(unknown_uploads, start=1):
+        src = getattr(f, "path", None) or getattr(f, "name", None) or str(f)
+        orig = (
+            getattr(f, "orig_name", None)
+            or getattr(f, "filename", None)
+            or os.path.basename(src)
+        )
+        lower = str(orig).lower()
+
+        # ZIP → extract FASTA files
+        if lower.endswith(".zip") or str(src).lower().endswith(".zip"):
+            _safe_extract_zip(src, unknown_dir)
+            continue
+
+        # FASTA → copy
+        if lower.endswith(FASTA_EXTS):
+            dst_name = os.path.basename(orig)
+        else:
+            # If Gradio provides a temp name without extension, keep it readable
+            dst_name = f"unknown_{idx}.fasta"
+
+        shutil.copy(src, os.path.join(unknown_dir, dst_name))
+
+
+def run_prediction(unknown_uploads, kmer_zip, seqtype, mode, identity, coverage, fdr):
+    if not unknown_uploads:
+        raise gr.Error("Please upload unknown FASTA files or a ZIP containing FASTA files.")
     if not kmer_zip:
         raise gr.Error("Please upload the k-mer results ZIP from Space 1.")
 
@@ -58,36 +101,23 @@ def run_prediction(unknown_files, kmer_zip, seqtype, mode, identity, coverage, f
     os.makedirs(unknown_dir, exist_ok=True)
     os.makedirs(outdir, exist_ok=True)
 
-    # Copy unknown FASTAs
-    for idx, f in enumerate(unknown_files, start=1):
-    src = getattr(f, "path", None) or getattr(f, "name", None) or str(f)
-    orig = (
-        getattr(f, "orig_name", None)
-        or getattr(f, "filename", None)
-        or os.path.basename(src)
-    )
-
-    lower = orig.lower()
-
-    # ZIP → extract FASTA files
-    if lower.endswith(".zip") or src.lower().endswith(".zip"):
-        _safe_extract_zip(src, unknown_dir)
-        continue
+    # Ingest unknown uploads (FASTA and/or ZIP)
+    _ingest_unknown_uploads(unknown_uploads, unknown_dir)
 
-    # FASTA file
-    if not lower.endswith((".fa", ".fasta", ".fas", ".fna")):
-        continue
-
-    dst = os.path.join(unknown_dir, os.path.basename(orig))
-    shutil.copy(src, dst)
-
-
-        shutil.copy(src, os.path.join(unknown_dir, os.path.basename(orig)))
+    # Ensure we actually got sequences
+    # (Lightweight check: presence of at least one fasta-like file)
+    found_any = any(
+        fn.lower().endswith(FASTA_EXTS)
+        for _, _, files in os.walk(unknown_dir)
+        for fn in files
+    )
+    if not found_any:
+        raise gr.Error("No FASTA files were found after processing your uploads. Please check your ZIP contents.")
 
     # K-mer ZIP path (ZIP-only)
     kmer_zip_path = getattr(kmer_zip, "path", None) or getattr(kmer_zip, "name", None) or str(kmer_zip)
     if not str(kmer_zip_path).lower().endswith(".zip"):
-        raise gr.Error("K-mer input must be a .zip file from Space 1.")
+        raise gr.Error("K-mer input must be a .zip file produced by Space 1.")
 
     # Run prediction
     kmer_predict.predict(
@@ -114,13 +144,13 @@ def run_prediction(unknown_files, kmer_zip, seqtype, mode, identity, coverage, f
 with gr.Blocks() as demo:
     gr.Markdown("# K-mer Sequence Predictor")
     gr.Markdown(
-        "Upload unknown FASTA sequences and the **kmer_results.zip** produced by the Unique k-mer Space."
+        "Upload **unknown sequences** (FASTA files or ZIP containing FASTA) and the **kmer_results.zip** from Space 1."
     )
 
-    unknown_files = gr.File(
-        label="Unknown FASTA files",
+    unknown_uploads = gr.File(
+        label="Unknown sequences (FASTA files or ZIP containing FASTA)",
         file_count="multiple",
-        file_types=[".fa", ".fasta", ".fas", ".fna"],
+        file_types=[".fa", ".fasta", ".fas", ".fna", ".zip"],
     )
 
     kmer_zip = gr.File(
@@ -146,7 +176,7 @@ with gr.Blocks() as demo:
 
     run_btn.click(
         fn=run_prediction,
-        inputs=[unknown_files, kmer_zip, seqtype, mode, identity, coverage, fdr],
+        inputs=[unknown_uploads, kmer_zip, seqtype, mode, identity, coverage, fdr],
         outputs=[out_plot, out_csv, out_zip],
     )