gistfile1.py

"""
# optional gff file as first arg.
>>> fi = FeatureIndexMemory('tests/data/grape.gff3')
>>> for i, (fid, feature) in enumerate(fi.iteritems()):
...    if i > 3: break
...    feature
FeatureNode('GSVIV00028344001', 3699817..3700731, strand='+', seqid='10')
FeatureNode('GSVIV00029007001', 9837193..9850255, strand='-', seqid='1')
FeatureNode('GSVIV00029792001', 6527551..6530322, strand='+', seqid='10')
FeatureNode('GSVIV00030315001', 2976811..2980733, strand='-', seqid='1')

>>> fi.seqids
['1', '10']


# by feature.attribs["ID"]
>>> fi['GSVIV00028344001']
FeatureNode('GSVIV00028344001', 3699817..3700731, strand='+', seqid='10')


#wrapper for get_features_for_range
>>> a = fi[3699817:3700731:'10'][0]
>>> b = fi.get_features_for_range(3699817, 3700731, '10')[0]
>>> a.start == b.start and a.end == b.end and a.attribs["ID"] == b.attribs["ID"]
True


>>> fi[3699817:3700731:'10']
[FeatureNode(start=3699817, end=3700731, seqid="10")]


# by seqid -- wrapper for get_features_for_seqid
>>> a = fi['10'][:2] 
>>> b = fi.get_features_for_seqid("10")[:2]
>>> [f.start for f in a] == [f.start for f in b]
True

"""
import gt
from gt import FeatureNode

__all__ = ['FeatureIndexMemory', 'FeatureNode']

class FeatureNode(gt.FeatureNode):
    def __repr__(self):
        return "FeatureNode('%s', %i..%i, strand='%s', seqid='%s')" % \
                (self.attribs["ID"], self.start, self.end, self.strand,
                self.seqid)

class FeatureIndexMemory(gt.FeatureIndexMemory):
    def __init__(self, gff3file=None, with_ids=True):
        gt.FeatureIndexMemory.__init__(self)
        if not gff3file is None:
            self.add_gff3file(gff3file)
            self.seqids = sorted(self.get_seqids())
        else:
            self.seqids = []
        self._ids = {}
        if with_ids:
            self._load_ids()

    def __getitem__(self, fid):
        if isinstance(fid, slice):
            return self.get_features_for_range(fid.start, fid.stop, str(fid.step))
        else:
            try:
                ptr = self._ids[fid]
                return FeatureNode.create_from_ptr(ptr)
            except KeyError:
                if fid in self.seqids:
                    return self.get_features_for_seqid(fid)
                else:
                    raise

    def _load_ids(self):
	self._ids = {}
        for seqid in self.seqids:
            for f in self.get_features_for_seqid(seqid):
                self._ids[f.attribs['ID']] = f._as_parameter_


    def iteritems(self):
        for fid, ptr in self._ids.iteritems():
            yield fid, FeatureNode.create_from_ptr(ptr)
        


if __name__ == "__main__":
    import doctest
    doctest.testmod()