r718: retrieve sequence from the index

python/README.rst

@@ -34,6 +34,8 @@ The following Python script demonstrates the key functionality of mappy:
 	import mappy as mp
 	a = mp.Aligner("test/MT-human.fa")  # load or build index
 	if not a: raise Exception("ERROR: failed to load/build index")
+	s = a.seq("MT_human", 100, 200)     # retrieve a subsequence from the index
+	print(mp.revcomp(s))                # reverse complement
 	for name, seq, qual in mp.fastx_read("test/MT-orang.fa"): # read a fasta/q sequence
 		for hit in a.map(seq): # traverse alignments
 			print("{}\t{}\t{}\t{}".format(hit.ctg, hit.r_st, hit.r_en, hit.cigar_str))
@@ -87,7 +89,15 @@ This method aligns :code:`seq` against the index. It is a generator, *yielding*
 a series of :code:`mappy.Alignment` objects. If :code:`seq2` is present, mappy
 performs paired-end alignment, assuming the two ends are in the FR orientation.
 Alignments of the two ends can be distinguished by the :code:`read_num` field
-(see below).
+(see Class mappy.Alignment below).
+
+.. code:: python
+
+	mappy.Aligner.seq(name, start=0, end=0x7fffffff)
+
+This method retrieves a (sub)sequence from the index and returns it as a Python
+string. :code:`None` is returned if :code:`name` is not present in the index or
+the start/end coordinates are invalid.
 
 Class mappy.Alignment
 ~~~~~~~~~~~~~~~~~~~~~
@@ -144,11 +154,12 @@ Miscellaneous Functions
 
 .. code:: python
 
-	mappy.fastx_read(fn)
+	mappy.fastx_read(fn, read_comment=False)
 
 This generator function opens a FASTA/FASTQ file and *yields* a
 :code:`(name,seq,qual)` tuple for each sequence entry. The input file may be
-optionally gzip'd.
+optionally gzip'd. If :code:`read_comment` is True, this generator yields
+a :code:`(name,seq,qual,comment)` tuple instead.
 
 .. code:: python
 

python/cmappy.h

@@ -112,4 +112,22 @@ static inline char *mappy_revcomp(int len, const uint8_t *seq)
 	return rev;
 }
 
+static char *mappy_fetch_seq(const mm_idx_t *mi, const char *name, int st, int en, int *len)
+{
+	int i, rid;
+	char *s;
+	*len = 0;
+	rid = mm_idx_name2id(mi, name);
+	if (rid < 0) return 0;
+	if (st >= mi->seq[i].len || st >= en) return 0;
+	if (en < 0 || en > mi->seq[i].len)
+		en = mi->seq[i].len;
+	s = (char*)malloc(en - st + 1);
+	*len = mm_idx_getseq(mi, rid, st, en, s);
+	for (i = 0; i < *len; ++i)
+		s[i] = "ACGTN"[(uint8_t)s[i]];
+	s[*len] = 0;
+	return s;
+}
+
 #endif

python/cmappy.pxd

@@ -58,7 +58,7 @@ cdef extern from "minimap.h":
 		mm_idx_seq_t *seq
 		uint32_t *S
 		mm_idx_bucket_t *B
-		void *km
+		void *km, *h
 
 	ctypedef struct mm_idx_reader_t:
 		pass
@@ -69,6 +69,8 @@ cdef extern from "minimap.h":
 	void mm_idx_destroy(mm_idx_t *mi)
 	void mm_mapopt_update(mm_mapopt_t *opt, const mm_idx_t *mi)
 
+	int mm_idx_index_name(mm_idx_t *mi)
+
 	#
 	# Mapping (key struct defined in cmappy.h below)
 	#
@@ -99,6 +101,7 @@ cdef extern from "cmappy.h":
 	void mm_reg2hitpy(const mm_idx_t *mi, mm_reg1_t *r, mm_hitpy_t *h)
 	void mm_free_reg1(mm_reg1_t *r)
 	mm_reg1_t *mm_map_aux(const mm_idx_t *mi, const char *seq1, const char *seq2, int *n_regs, mm_tbuf_t *b, const mm_mapopt_t *opt)
+	char *mappy_fetch_seq(const mm_idx_t *mi, const char *name, int st, int en, int *l)
 
 	ctypedef struct kstring_t:
 		unsigned l, m

python/mappy.pyx

@@ -123,6 +123,7 @@ cdef class Aligner:
 			self._idx = cmappy.mm_idx_reader_read(r, n_threads) # NB: ONLY read the first part
 			cmappy.mm_idx_reader_close(r)
 			cmappy.mm_mapopt_update(&self.map_opt, self._idx)
+			cmappy.mm_idx_index_name(self._idx)
 
 	def __dealloc__(self):
 		if self._idx is not NULL:
@@ -140,8 +141,13 @@ cdef class Aligner:
 		if self._idx is NULL: return None
 		if buf is None: b = ThreadBuffer()
 		else: b = buf
-		if seq2 is None: regs = cmappy.mm_map_aux(self._idx, str.encode(seq), NULL, &n_regs, b._b, &self.map_opt)
-		else: regs = cmappy.mm_map_aux(self._idx, str.encode(seq), str.encode(seq2), &n_regs, b._b, &self.map_opt)
+
+		_seq = seq if isinstance(seq, bytes) else seq.encode()
+		if seq2 is None:
+			regs = cmappy.mm_map_aux(self._idx, _seq, NULL,  &n_regs, b._b, &self.map_opt)
+		else:
+			_seq2 = seq2 if isinstance(seq2, bytes) else seq2.encode()
+			regs = cmappy.mm_map_aux(self._idx, _seq, _seq2, &n_regs, b._b, &self.map_opt)
 
 		for i in range(n_regs):
 			cmappy.mm_reg2hitpy(self._idx, &regs[i], &h)
@@ -153,7 +159,24 @@ cdef class Aligner:
 			cmappy.mm_free_reg1(&regs[i])
 		free(regs)
 
-def fastx_read(fn):
+	def seq(self, str name, int start=0, int end=0x7fffffff):
+		cdef int l
+		cdef char *s = cmappy.mappy_fetch_seq(self._idx, name.encode(), start, end, &l)
+		if l == 0: return None
+		r = s[:l] if isinstance(s, str) else s[:l].decode()
+		free(s)
+		return r
+
+	@property
+	def k(self): return self._idx.k
+
+	@property
+	def w(self): return self._idx.w
+
+	@property
+	def n_seq(self): return self._idx.n_seq
+
+def fastx_read(fn, read_comment=False):
 	cdef cmappy.kseq_t *ks
 	ks = cmappy.mm_fastx_open(str.encode(fn))
 	if ks is NULL: return None
@@ -162,7 +185,12 @@ def fastx_read(fn):
 		else: qual = None
 		name = ks.name.s if isinstance(ks.name.s, str) else ks.name.s.decode()
 		seq = ks.seq.s if isinstance(ks.seq.s, str) else ks.seq.s.decode()
-		yield name, seq, qual
+		if read_comment:
+			if ks.comment.l > 0: comment = ks.comment.s if isinstance(ks.comment.s, str) else ks.comment.s.decode()
+			else: comment = None
+			yield name, seq, qual, comment
+		else:
+			yield name, seq, qual
 	cmappy.mm_fastx_close(ks)
 
 def revcomp(seq):