Implement KmerTable in Rust

src/biotite/sequence/align/buckets.py

@@ -6,10 +6,13 @@
 __author__ = "Patrick Kunzmann"
 __all__ = ["bucket_number"]
 
-from os.path import dirname, join, realpath
+import functools
+from pathlib import Path
+from typing import Any
 import numpy as np
+from biotite.typing import NDArray1
 
-_primes = None
+_PRIMES_PATH = Path(__file__).parent / "primes.txt"
 
 
 def bucket_number(n_kmers: int, load_factor: float = 0.8) -> int:
@@ -53,19 +56,22 @@ def bucket_number(n_kmers: int, load_factor: float = 0.8) -> int:
     Hence, all *k-mers* with the same suffix would be stored in the same
     bin.
     """
-    global _primes
-    if _primes is None:
-        with open(join(dirname(realpath(__file__)), "primes.txt")) as file:
-            _primes = np.array(
-                [
-                    int(line)
-                    for line in file.read().splitlines()
-                    if len(line) != 0 and line[0] != "#"
-                ]
-            )
-
+    primes = _get_primes()
     number = int(n_kmers / load_factor)
-    index = np.searchsorted(_primes, number, side="left")
-    if index == len(_primes):
+    index = np.searchsorted(primes, number, side="left")
+    if index == len(primes):
         raise ValueError("Number of buckets too large")
-    return _primes[index]
+    return primes[index]
+
+
+@functools.cache
+def _get_primes() -> NDArray1[Any, np.integer]:
+    with open(_PRIMES_PATH) as file:
+        _primes = np.array(
+            [
+                int(line)
+                for line in file.read().splitlines()
+                if len(line) != 0 and line[0] != "#"
+            ]
+        )
+    return _primes

src/biotite/sequence/align/kmertable.py

@@ -0,0 +1,16 @@
+# This source code is part of the Biotite package and is distributed
+# under the 3-Clause BSD License. Please see 'LICENSE.rst' for further
+# information.
+
+from __future__ import annotations
+
+__name__ = "biotite.sequence.align"
+__author__ = "Patrick Kunzmann"
+__all__ = ["KmerTable", "BucketKmerTable"]
+
+from biotite.rust.sequence.align import BucketKmerTable, KmerTable
+
+# The classes are implemented in Rust and therefore are not generic at runtime;
+# allow ``KmerTable[S]`` subscription so the annotated stub works
+KmerTable.__class_getitem__ = classmethod(lambda cls, params: cls)
+BucketKmerTable.__class_getitem__ = classmethod(lambda cls, params: cls)

src/biotite/sequence/align/kmertable.pyi

@@ -1,7 +1,7 @@
 __all__ = ["KmerTable", "BucketKmerTable"]
 
 from collections.abc import Iterable, Iterator
-from typing import Any, Generic
+from typing import Generic
 import numpy as np
 from biotite.sequence.align.kmeralphabet import KmerAlphabet
 from biotite.sequence.align.kmersimilarity import SimilarityRule
@@ -10,7 +10,6 @@ from biotite.sequence.sequence import Sequence
 from biotite.typing import C2, C3, C4, M, N, NDArray1, NDArray2, S
 
 class KmerTable(Generic[S]):
-    def __init__(self, kmer_alphabet: KmerAlphabet[S]) -> None: ...
     @property
     def kmer_alphabet(self) -> KmerAlphabet[S]: ...
     @property
@@ -74,12 +73,8 @@ class KmerTable(Generic[S]):
     def __reversed__(self) -> Iterator[int]: ...
     def __eq__(self, item: object) -> bool: ...
     def __str__(self) -> str: ...
-    def __getnewargs_ex__(self) -> tuple[tuple, dict]: ...
-    def __getstate__(self) -> Any: ...
-    def __setstate__(self, state: Any) -> None: ...
 
 class BucketKmerTable(Generic[S]):
-    def __init__(self, n_buckets: int, kmer_alphabet: KmerAlphabet[S]) -> None: ...
     @property
     def kmer_alphabet(self) -> KmerAlphabet[S]: ...
     @property
@@ -140,6 +135,3 @@ class BucketKmerTable(Generic[S]):
     def __len__(self) -> int: ...
     def __eq__(self, item: object) -> bool: ...
     def __str__(self) -> str: ...
-    def __getnewargs_ex__(self) -> tuple[tuple, dict]: ...
-    def __getstate__(self) -> Any: ...
-    def __setstate__(self, state: Any) -> None: ...