indexing refactor

src/beignet/datasets/__uni_ref_dataset.py

@@ -1,4 +1,3 @@
-import re
 from os import PathLike
 from pathlib import Path
 from typing import Callable
@@ -51,8 +50,6 @@ def __init__(
 
         name = self.__class__.__name__.replace("Dataset", "")
 
-        self._pattern = re.compile(r"^UniRef.+_([A-Z0-9]+)\s.+$")
-
         super().__init__(
             pooch.retrieve(
                 url,
@@ -71,14 +68,12 @@ def __init__(
         self.target_transform = target_transform
 
     def __getitem__(self, index: int) -> (str, str):
-        target, sequence = self.get(index)
-
-        (target,) = re.search(self._pattern, target).groups()
+        input, target = self.get(index)
 
         if self.transform:
-            sequence = self.transform(sequence)
+            input = self.transform(input)
 
         if self.target_transform:
             target = self.target_transform(target)
 
-        return sequence, target
+        return input, target

src/beignet/datasets/_fasta_dataset.py

@@ -1,9 +1,9 @@
-import subprocess
 from os import PathLike
 from pathlib import Path
 from typing import Callable, Tuple, TypeVar
 
 import numpy
+import tqdm
 
 from beignet.io import ThreadSafeFile
 
@@ -37,10 +37,7 @@ def __init__(
         else:
             self.offsets, sizes = self._build_index()
 
-            numpy.save(
-                f"{offsets}",
-                numpy.stack([self.offsets, sizes]),
-            )
+            numpy.save(f"{offsets}", numpy.stack([self.offsets, sizes]))
 
         self.transform = transform
 
@@ -57,44 +54,50 @@ def __getitem__(self, index: int) -> Tuple[str, str]:
     def __len__(self) -> int:
         return self.offsets.size
 
-    def get(self, index: int) -> str:
+    def get(self, index: int) -> (str, str):
         self.data.seek(self.offsets[index])
 
         if index == len(self) - 1:
             data = self.data.read()
         else:
-            data = self.data.read(
-                self.offsets[index + 1] - self.offsets[index],
-            )
+            data = self.data.read(self.offsets[index + 1] - self.offsets[index])
 
         description, *sequence = data.split("\n")
 
-        return "".join(sequence)
-
-    def _build_index(self) -> tuple[numpy.ndarray, numpy.ndarray]:
-        # TODO: rewrite in Rust (using `libripgrep`) or similar to remove
-        #  dependency on `grep` and `awk`. — Allen (Tuesday, November 29, 2022)
-        return (
-            numpy.fromstring(
-                subprocess.check_output(
-                    f"cat {self.root} "
-                    f"| tqdm --bytes --total $(wc -c < {self.root})"
-                    "| grep --byte-offset '^>' -o | cut -d: -f1",
-                    shell=True,
-                ),
-                dtype=numpy.int64,
-                sep=" ",
-            ),
-            numpy.fromstring(
-                subprocess.check_output(
-                    f"cat {self.root} "
-                    f"| tqdm --bytes --total $(wc -c < {self.root})"
-                    '| awk \'/^>/ {print "";next;} { printf("%s",$0);}\' '
-                    "| tail -n+2 | awk "
-                    "'{print length($1)}'",
-                    shell=True,
-                ),
-                dtype=numpy.int64,
-                sep=" ",
-            ),
-        )
+        return "".join(sequence), description
+
+    def _build_index(self) -> (numpy.ndarray, numpy.ndarray):
+        with open(self.root, "r") as file:
+            content = file.read()
+
+        offsets, sizes = [], []
+
+        current_offset, current_size = 0, 0
+
+        parsing = False
+
+        for sequence in tqdm.tqdm(content.splitlines(keepends=True)):
+            characters = len(sequence)
+
+            if sequence.startswith(">"):
+                if parsing:
+                    sizes = [*sizes, current_size]
+
+                    current_size = 0
+
+                offsets = [*offsets, current_offset]
+
+                parsing = True
+            elif parsing:
+                current_size = current_size + len(sequence.rstrip("\n"))
+
+            current_offset = current_offset + characters
+
+        if parsing:
+            sizes = [*sizes, current_size]
+
+        offsets = numpy.array(offsets, dtype=numpy.int64)
+
+        sizes = numpy.array(sizes, dtype=numpy.int64)
+
+        return offsets, sizes