taxo.pyx 12.2 KB
Newer Older
1 2
#cython: language_level=3

3
from obitools3.utils cimport str2bytes, bytes2str, tobytes, tostr
4

5 6
from ..capi.obitaxonomy cimport obi_taxonomy_exists, \
                                obi_read_taxonomy, \
7 8 9 10 11 12
                                obi_read_taxdump, \
                                obi_write_taxonomy, \
                                obi_close_taxonomy, \
                                obi_taxo_get_taxon_with_taxid, \
                                obi_taxo_add_local_taxon, \
                                obi_taxo_add_preferred_name_with_taxon, \
13 14 15 16
                                obi_taxo_rank_index_to_label, \
                                obi_taxo_get_species, \
                                obi_taxo_get_genus, \
                                obi_taxo_get_family, \
17 18 19 20
                                ecotx_t                
                                           
from cpython.pycapsule cimport PyCapsule_New, PyCapsule_GetPointer

21 22
import tarfile

23

24
cdef class Taxonomy(OBIWrapper) :
25
    # TODO function to import taxonomy?
26 27 28 29 30
    
    cdef inline OBIDMS_taxonomy_p pointer(self) :
        return <OBIDMS_taxonomy_p>(self._pointer)        


31 32 33 34 35 36 37 38 39 40
    @staticmethod
    def exists(DMS dms, object name) :
        e = obi_taxonomy_exists(dms.pointer(), tobytes(name))
        if e < 0:
            raise RuntimeError("Error : Cannot check if taxonomy %s exists" 
                               % tostr(name))
        else:
            return e
        

41
    @staticmethod
42
    def open(DMS dms, object name) :
43
        
44 45
        cdef void*    pointer
        cdef Taxonomy taxo
46 47 48 49 50 51 52 53 54 55 56 57
                
        pointer = <void*>obi_read_taxonomy(dms.pointer(), tobytes(name), True)        
        if pointer == NULL :
            raise RuntimeError("Error : Cannot read taxonomy %s" 
                               % tostr(name))

        taxo = OBIWrapper.new_wrapper(Taxonomy, pointer)

        dms.register(taxo)

        taxo._dms = dms
        taxo._name = tobytes(name)
58
        
59 60 61
        taxo._ranks = []
        for r in range((<OBIDMS_taxonomy_p>pointer).ranks.count) :
            taxo._ranks.append(obi_taxo_rank_index_to_label(r, (<OBIDMS_taxonomy_p>pointer).ranks))
62

63 64 65 66
        return taxo


    @staticmethod
67
    def open_taxdump(DMS dms, object path) :
68
        
69 70
        cdef void*    pointer
        cdef Taxonomy taxo
71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90
        cdef bytes    path_b
        cdef int      idx
        
        path_b = tobytes(path)
        folder_path = path_b

        if path_b.endswith(b"tar.gz") or path_b.endswith(b"tar"):
            idx = path_b.index(b".tar")
            folder_path = path_b[:idx]

        if path_b.endswith(b"tar.gz"):
            tar = tarfile.open(path_b, "r:gz")
            tar.extractall(path=tostr(folder_path))
            tar.close()
        elif path_b.endswith(b"tar"):
            tar = tarfile.open(path_b, "r:")
            tar.extractall(path=tostr(folder_path))
            tar.close()
            
        pointer = <void*>obi_read_taxdump(folder_path)        
91 92
        if pointer == NULL :
            raise RuntimeError("Error : Cannot read taxonomy %s" 
93
                               % tostr(folder_path))
94
        
95
        taxo = OBIWrapper.new_wrapper(Taxonomy, pointer)
96 97 98 99

        dms.register(taxo)

        taxo._dms = dms
100
        taxo._name = folder_path
101 102 103 104
        
        taxo._ranks = []
        for r in range((<OBIDMS_taxonomy_p>pointer).ranks.count) :
            taxo._ranks.append(obi_taxo_rank_index_to_label(r, (<OBIDMS_taxonomy_p>pointer).ranks))
105
                
106 107
        return taxo

108
    
109
    def __getitem__(self, object ref):     
110
        if type(ref) == int :
111
            return self.get_taxon_by_taxid(ref)
112
        else :
Celine Mercier's avatar
Celine Mercier committed
113
            raise NotImplementedError()
114 115


116
    cpdef Taxon get_taxon_by_taxid(self, int taxid):   # TODO check deleted taxon behavior (403122)
117 118 119 120 121 122 123 124 125 126
        cdef ecotx_t* taxon_p
        cdef object   taxon_capsule        
        taxon_p = obi_taxo_get_taxon_with_taxid(self.pointer(), taxid)
        if taxon_p == NULL:
            raise Exception("Error getting a taxon with given taxid", taxid)
        taxon_capsule = PyCapsule_New(taxon_p, NULL, NULL)
        return Taxon(taxon_capsule, self)


    cpdef Taxon get_taxon_by_idx(self, int idx):
127 128 129
        cdef ecotx_t* taxa
        cdef ecotx_t* taxon_p
        cdef object   taxon_capsule
130 131
        if idx >= self.pointer().taxa.count :
            raise Exception("Error getting a taxon with given index: no taxid with this index", idx)
132 133 134
        taxa = self.pointer().taxa.taxon
        taxon_p = <ecotx_t*> (taxa+idx)
        taxon_capsule = PyCapsule_New(taxon_p, NULL, NULL)
135
        return Taxon(taxon_capsule, self)
136

137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190

    cpdef object get_species(self, int taxid):
        cdef ecotx_t* taxon_p
        cdef ecotx_t* species_p
        taxon_p = obi_taxo_get_taxon_with_taxid(self.pointer(), taxid)
        if taxon_p == NULL:
            raise Exception("Error getting a taxon with given taxid", taxid)
        species_p = obi_taxo_get_species(taxon_p, self.pointer())
        if species_p == NULL :
            return None
        else :
            return <int>(species_p.taxid)


    cpdef object get_genus(self, int taxid):
        cdef ecotx_t* taxon_p
        cdef ecotx_t* genus_p
        taxon_p = obi_taxo_get_taxon_with_taxid(self.pointer(), taxid)
        if taxon_p == NULL:
            raise Exception("Error getting a taxon with given taxid", taxid)
        genus_p = obi_taxo_get_genus(taxon_p, self.pointer())
        if genus_p == NULL :
            return None
        else :
            return <int>(genus_p.taxid)


    cpdef object get_family(self, int taxid):
        cdef ecotx_t* taxon_p
        cdef ecotx_t* family_p
        taxon_p = obi_taxo_get_taxon_with_taxid(self.pointer(), taxid)
        if taxon_p == NULL:
            raise Exception("Error getting a taxon with given taxid", taxid)
        family_p = obi_taxo_get_family(taxon_p, self.pointer())
        if family_p == NULL :
            return None
        else :
            return <int>(family_p.taxid)


    cpdef bytes get_scientific_name(self, int taxid):
        cdef ecotx_t* taxon_p
        taxon_p = obi_taxo_get_taxon_with_taxid(self.pointer(), taxid)
        if taxon_p == NULL:
            raise Exception("Error getting a taxon with given taxid", taxid)
        return taxon_p.name
    
    
    cpdef bytes get_rank(self, int taxid):
        cdef ecotx_t* taxon_p
        taxon_p = obi_taxo_get_taxon_with_taxid(self.pointer(), taxid)
        if taxon_p == NULL:
            raise Exception("Error getting a taxon with given taxid", taxid)
        return self._ranks[taxon_p.rank]
191 192 193 194 195 196 197 198 199 200


    cpdef object get_taxon_at_rank(self, int taxid, object rank):
        if isinstance(rank, str) or isinstance(rank, bytes):
            rank = self._ranks.index(tobytes(rank))
        try:
            return [x.taxid for x in self.parental_tree_iterator(taxid) if x.rank_idx==rank][0]
        except IndexError:
            return None

201 202 203 204 205
    
    def __len__(self):
        return self.pointer().taxa.count


206 207 208 209 210 211 212
    def __iter__(self):
         
        cdef ecotx_t* taxa
        cdef ecotx_t* taxon_p
        cdef object   taxon_capsule
        cdef int      t
 
213
        taxa = self.pointer().taxa.taxon
214 215
 
        # Yield each taxid
216
        for t in range(self.pointer().taxa.count):
217 218
            taxon_p = <ecotx_t*> (taxa+t)
            taxon_capsule = PyCapsule_New(taxon_p, NULL, NULL)
219
            yield Taxon(taxon_capsule, self)
220 221


222
    cpdef write(self, object prefix) :
223
        if obi_write_taxonomy(self._dms.pointer(), self.pointer(), tobytes(prefix)) < 0 :
224 225 226 227 228
            raise Exception("Error writing the taxonomy to binary files")
    
    
    cpdef int add_taxon(self, str name, str rank_name, int parent_taxid, int min_taxid=10000000) :
        cdef int taxid
229
        taxid = obi_taxo_add_local_taxon(self.pointer(), tobytes(name), tobytes(rank_name), parent_taxid, min_taxid)
230 231 232 233
        if taxid < 0 :
            raise Exception("Error adding a new taxon to the taxonomy")
        else :
            return taxid
234 235


236
    def close(self) :        
237 238 239 240 241 242 243 244 245 246 247 248
        if self.active() :
            self._dms.unregister(self)
            OBIWrapper.close(self)      
            if (obi_close_taxonomy(self.pointer()) < 0) :
                raise Exception("Problem closing the taxonomy %s" % 
                                self._name)


    # name property getter
    @property
    def name(self):
        return self._name
249 250 251 252 253 254 255 256
    
    
    def parental_tree_iterator(self, int taxid):
        """
           return parental tree for given taxonomic id starting from
           first ancestor to the root.
        """
        cdef Taxon taxon
257 258 259 260
        try:
            taxon = self.get_taxon_by_taxid(taxid)
        except:  # TODO error handling (related to deleted taxon thing)
            raise StopIteration
261 262 263 264 265 266 267 268 269
        if taxon is not None:
            while taxon.parent.taxid != 1:   # TODO was 0 before? 
                yield taxon
                taxon = taxon.parent
            yield self[1]
        else:
            raise StopIteration


270 271 272 273
    def is_ancestor(self, int ancestor_taxid, int taxid):
        return ancestor_taxid in [x.taxid for x in self.parental_tree_iterator(taxid)]


274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306
    def last_common_taxon(self, *taxids):
        
        cdef list  t1
        cdef list  t2
        cdef Taxon x
        cdef int   count
        cdef int   i
        cdef int   ancestor
        
        if not taxids:
            return None
        if len(taxids)==1:
            return taxids[0]
                
        if len(taxids)==2:
            t1 = [x.taxid for x in self.parental_tree_iterator(taxids[0])]
            t2 = [x.taxid for x in self.parental_tree_iterator(taxids[1])]
            t1.reverse()
            t2.reverse()
            
            count = min(len(t1),len(t2))
            i=0
            while(i < count and t1[i]==t2[i]):
                i+=1
            i-=1
                        
            return t1[i]
        
        ancestor = taxids[0]
        for taxon in taxids[1:]:
            ancestor = self.last_common_taxon(ancestor, taxon)

        return ancestor
307

308
    
309
cdef class Taxon :    # TODO dict subclass?
310
    
311
    def __init__(self, object taxon_capsule, Taxonomy tax) :
312 313 314 315 316 317
        self._pointer = <ecotx_t*> PyCapsule_GetPointer(taxon_capsule, NULL)
        if self._pointer == NULL :
            raise Exception("Error reading a taxon (NULL pointer)")
        self._tax = tax
    
    
318
    # To test equality
319
    def __richcmp__(self, Taxon taxon2, int op):
320 321
        return (self.name == taxon2.name) and \
               (self.taxid == taxon2.taxid) and \
322
               (self.rank_idx == taxon2.rank_idx) and \
323 324 325
               (self.farest == taxon2.farest) and \
               (self.parent.taxid == taxon2.parent.taxid) and \
               (self.preferred_name == taxon2.preferred_name)
326 327


328 329 330
    # name property getter
    @property
    def name(self):
331
        return self._pointer.name
332

333 334 335 336
    # taxid property getter
    @property
    def taxid(self):
        return self._pointer.taxid
337

338 339 340
    # rank property getter
    @property
    def rank(self):
341
        return ((self._tax)._ranks)[(self._pointer).rank]
342

343 344 345 346 347
    # rank_idx property getter
    @property
    def rank_idx(self):
        return (self._pointer).rank

348 349 350 351
    # farest property getter
    @property
    def farest(self):
        return self._pointer.farest
352

353 354 355 356 357
    # parent property getter
    @property
    def parent(self):
        cdef object parent_capsule
        parent_capsule = PyCapsule_New(self._pointer.parent, NULL, NULL)
358
        return Taxon(parent_capsule, self._tax)
359

360 361 362 363 364
    # preferred name property getter and setter
    @property
    def preferred_name(self):
        if self._pointer.preferred_name != NULL :
            return bytes2str(self._pointer.preferred_name)
365

366 367
    @preferred_name.setter
    def preferred_name(self, str new_preferred_name) :  # @DuplicatedSignature
368
        if (obi_taxo_add_preferred_name_with_taxon(self._tax.pointer(), self._pointer, tobytes(new_preferred_name)) < 0) :
369
            raise Exception("Error adding a new preferred name to a taxon")
370

371 372 373 374
    def __repr__(self):
        d = {}
        d['taxid']          = self.taxid
        d['name']           = self.name
375
        d['rank']           = self.rank
376
        d['rank_idx']       = self.rank_idx
377 378 379 380 381 382
        d['preferred name'] = self.preferred_name
        d['parent']         = self.parent.taxid
        d['farest']         = self.farest
        return str(d)