import ctypes from pathlib import Path _suffixtool = ctypes.CDLL("./suffixtool.so") _suffixtool.csa_build.argtypes = [ctypes.c_char_p, ctypes.c_char_p] _suffixtool.csa_build.restype = ctypes.c_int _suffixtool.csa_load.argtypes = [ctypes.c_char_p] _suffixtool.csa_load.restype = ctypes.c_void_p _suffixtool.csa_free.argtypes = [ctypes.c_void_p] _suffixtool.csa_free.restype = None _suffixtool.csa_count.argtypes = [ctypes.c_void_p, ctypes.c_char_p] _suffixtool.csa_count.restype = ctypes.c_ssize_t def _path_to_bytes(x): if isinstance(x, Path): return bytes(x) if isinstance(x, str): return x.encode('utf-8') if isinstance(x, bytes): return x raise RuntimeError(f'Expected path-like object but got {x}') class CSA: def __init__(self, ptr): self.ptr = ptr @staticmethod def build(input_file, output_file): res = _suffixtool.csa_build(_path_to_bytes(input_file), _path_to_bytes(output_file)) if not res: raise RuntimeError('Failed to build CSA') @classmethod def load(cls, input_file): res = _suffixtool.csa_load(_path_to_bytes(input_file)) if res is None: raise RuntimeError('Failed to load CSA') return cls(res) def count(self, byte_seq): if self.ptr is None: raise RuntimeError('Called CSA.count() but CSA was already deleted') res = _suffixtool.csa_count(self.ptr, byte_seq) if res == -1: raise RuntimeError('Failed to count CSA byte sequence occurrences') return res def delete_now(self): if self.ptr is not None: _suffixtool.csa_free(self.ptr) self.ptr = None def __del__(self): self.delete_now()