Add inital wrapperHEADmain

2 files changed, 113 insertions, 0 deletions

diff --git a/pyqidx/__init__.py b/pyqidx/__init__.py
new file mode 100644
index 0000000..3b515a1
--- /dev/null
+++ b/pyqidx/__init__.py
@@ -0,0 +1,5 @@
+from pyqidx.qidx import (
+	AlignmentRecord,
+	QidxException,
+	QueryIndexFile
+)
diff --git a/pyqidx/qidx.py b/pyqidx/qidx.py
new file mode 100644
index 0000000..95f5113
--- /dev/null
+++ b/pyqidx/qidx.py
@@ -0,0 +1,108 @@
+from pyqidx._qidx import ffi, lib
+
+from typing import (
+	NamedTuple,
+	Generator,
+	BinaryIO,
+	Union,
+	Optional
+)
+
+class _QidxError( object ):
+
+	code : int
+
+	def __init__( self, code : int ):
+		self.code = code
+
+	@staticmethod
+	def from_code( code ):
+		return _QidxError( code )
+
+	@property
+	def message( self ):
+		_msg = ffi.string( lib.qidx_strerr( self.code ) )
+		return str( _msg , encoding='UTF-8' )
+
+class QidxException( Exception ):
+	
+	@staticmethod
+	def from_code( err_code ):
+		return QidxException( _QidxError( err_code ).message )
+
+class AlignmentRecord( NamedTuple ):
+	tid : int
+	pos : int
+	vptr : int
+
+class _Qidx( object ):
+
+	def __init__( self, fd : int ):
+		self._qidx = ffi.new('qidx_fp_t **')
+
+		err = lib.qidx_open( self._qidx, fd )
+
+		if err != lib.QIDX_OK:
+			raise QidxException.from_code( err )
+
+	def lookup_alignments( self, query_name : str ) \
+		-> Generator[AlignmentRecord, None, None]:
+		qname = ffi.new("char[]", bytes( query_name, encoding="UTF-8" ))
+
+		rec = ffi.new('qidx_record_t **')
+		err = lib.qidx_lookup_alnrec(self._qidx[0], qname, rec)
+
+		if not rec[0]:
+			return
+
+		for i in range(rec[0].n_alns):
+			aln = rec[0].alns[i]
+			yield AlignmentRecord( aln.tid, aln.pos, aln.vptr )
+
+		lib.qidx_free_alnrec( rec[0] )
+
+	def close( self ):
+		err = lib.qidx_close( self._qidx[0] )
+
+		if err != lib.QIDX_OK:
+			raise QidxException.from_code( err )
+
+class QueryIndexFile( object ):
+
+	_filename : Optional[ str ]
+	_file : BinaryIO
+	_qidx : _Qidx
+
+	def __init__( self, filename_or_file : Union[ BinaryIO, str ] ):
+
+		if isinstance( filename_or_file, str ):
+			self._filename = filename_or_file
+			self._file = None
+		else:
+			self._file = filename_or_file
+			self._filename = None
+
+		self._qidx = None
+
+	def open( self ):
+		if self._filename:
+			self._file = open( self._filename, 'rb' )
+
+		self._qidx = _Qidx( self._file.fileno() )
+
+	def close( self ):
+		self._qidx.close()
+
+		if self._filename:
+			self._file.close()
+
+	def __enter__( self ):
+		self.open()
+		return self
+
+	def __exit__( self, exc_type, exc_value, tb ):
+		self.close()
+
+	def lookup_alignments( self, query_name : str ) \
+		-> Generator[AlignmentRecord, None, None]:
+		return (yield from self._qidx.lookup_alignments( query_name ))