From de52799d94be8f2aac96942dca6a1d995b337bc0 Mon Sep 17 00:00:00 2001
From: jackh726 <git@jackhuey.me>
Date: Tue, 13 Aug 2024 17:34:27 +0000
Subject: [PATCH] When opening a bigBed with a file-like, use GenericBBIFile,
 otherwise we read initial bytes twice

---
 pybigtools/src/lib.rs        | 28 +++++++++++++---------------
 pybigtools/tests/test_api.py |  4 ++++
 2 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/pybigtools/src/lib.rs b/pybigtools/src/lib.rs
index a7a2ca2..53a25b6 100644
--- a/pybigtools/src/lib.rs
+++ b/pybigtools/src/lib.rs
@@ -17,7 +17,7 @@ use bigtools::utils::misc::{
 use bigtools::{
     BBIFileRead, BBIReadError as _BBIReadError, BedEntry, BigBedRead as BigBedReadRaw,
     BigBedWrite as BigBedWriteRaw, BigWigRead as BigWigReadRaw, BigWigWrite as BigWigWriteRaw,
-    CachedBBIFileRead, Value, ZoomRecord,
+    CachedBBIFileRead, GenericBBIRead, Value, ZoomRecord,
 };
 
 use bigtools::utils::reopen::Reopen;
@@ -2385,22 +2385,20 @@ fn open(py: Python, path_url_or_file_like: PyObject, mode: Option<String>) -> Py
             "Unknown argument for `path_url_or_file_like`. Not a file path string or url, and not a file-like object.",
         ))),
     };
-    let read = match BigWigReadRaw::open(file_like.clone()) {
-        Ok(bwr) => BBIRead {
-            bbi: BBIReadRaw::BigWigFileLike(bwr.cached()),
+    let read = match GenericBBIRead::open(file_like.clone()) {
+        Ok(GenericBBIRead::BigWig(bigwig)) => BBIRead {
+            bbi: BBIReadRaw::BigWigFileLike(bigwig.cached()),
         }
         .into_py(py),
-        Err(_) => match BigBedReadRaw::open(file_like) {
-            Ok(bbr) => BBIRead {
-                bbi: BBIReadRaw::BigBedFileLike(bbr.cached()),
-            }
-            .into_py(py),
-            Err(e) => {
-                return Err(PyErr::new::<BBIReadError, _>(format!(
-                    "File-like object is not a bigWig or bigBed. Or there was just a problem reading: {e}",
-                )))
-            }
-        },
+        Ok(GenericBBIRead::BigBed(bigbed)) => BBIRead {
+            bbi: BBIReadRaw::BigBedFileLike(bigbed.cached()),
+        }
+        .into_py(py),
+        Err(e) => {
+            return Err(PyErr::new::<BBIReadError, _>(format!(
+            "File-like object is not a bigWig or bigBed. Or there was just a problem reading: {e}",
+        )))
+        }
     };
     Ok(read)
 }
diff --git a/pybigtools/tests/test_api.py b/pybigtools/tests/test_api.py
index 047e3cb..017abbd 100644
--- a/pybigtools/tests/test_api.py
+++ b/pybigtools/tests/test_api.py
@@ -46,6 +46,10 @@ def test_open_filelike():
         with pybigtools.open(f, "r") as b:
             assert b.chroms() == {"chr17": 83_257_441}
 
+    with open(TEST_DIR / "data/bigBedExample.bb", "rb") as f:
+        with pybigtools.open(f, "r") as b:
+            assert b.chroms("chr21") == 48_129_895
+
     # BytesIO
     with open(REPO_ROOT / "bigtools/resources/test/valid.bigWig", "rb") as f:
         bw_bytes = f.read()