_obitaxo.pyx 3.76 KB
Newer Older
1 2 3 4 5
#cython: language_level=3

from obitools3.utils   cimport bytes2str, str2bytes

from .capi.obitaxonomy cimport obi_read_taxonomy, \
6 7
                               obi_read_taxdump, \
                               obi_write_taxonomy, \
8
                               obi_close_taxonomy, \
9
                               obi_taxo_get_taxon_with_taxid, \
10
                               obi_taxonomy_add_local_taxon, \
11
                               ecotx_t                            
12

13 14 15 16

from ._obidms cimport OBIDMS

from cpython.pycapsule cimport PyCapsule_New, PyCapsule_GetPointer
17
from logging import raiseExceptions
18 19 20 21


cdef class OBI_Taxonomy :

22 23
    # TODO function to import taxonomy?

24
    def __init__(self, OBIDMS dms, str name, bint taxdump=False) :
25
        
26
        self._dms = dms
27
        self._name = name
28 29 30 31
        if taxdump :
            self._pointer = obi_read_taxdump(str2bytes(name))
        else :
            self._pointer = obi_read_taxonomy(dms._pointer, str2bytes(name), True)    # TODO discuss
32 33
        # TODO if not found in DMS, try to import?
        
34 35 36 37
    
    def __getitem__(self, object ref):
        
        cdef ecotx_t* taxon_p
38
        cdef object   taxon_capsule
39 40
        
        if type(ref) == int :
41
            taxon_p = obi_taxo_get_taxon_with_taxid(self._pointer, ref)
42 43
            if taxon_p == NULL :
                raise Exception("Taxon not found")
44 45
            taxon_capsule = PyCapsule_New(taxon_p, NULL, NULL)
            return OBI_Taxon(taxon_capsule)
46 47 48
        else :
            raise Exception("Not implemented")
        
49
    
50 51 52 53 54 55 56 57 58 59 60
    def __iter__(self):
         
        cdef ecotx_t* taxa
        cdef ecotx_t* taxon_p
        cdef object   taxon_capsule
        cdef int      t
 
        taxa = self._pointer.taxa.taxon
 
        # Yield each taxid
        for t in range(self._pointer.taxa.count):
61
            taxon_p = <ecotx_t*> (taxa+t)
62 63
            taxon_capsule = PyCapsule_New(taxon_p, NULL, NULL)
            yield OBI_Taxon(taxon_capsule)
64 65 66 67 68 69


    cpdef write(self, str prefix) :
        if obi_write_taxonomy(self._dms._pointer, self._pointer, str2bytes(prefix)) < 0 :
            raise Exception("Error writing the taxonomy to binary files")

70 71 72 73 74 75 76 77 78
    
    cpdef int add_taxon(self, str name, str rank_name, int parent_taxid, int min_taxid=10000000) :
        cdef int taxid
        taxid = obi_taxonomy_add_local_taxon(self._pointer, str2bytes(name), str2bytes(rank_name), parent_taxid, min_taxid)
        if taxid < 0 :
            raise Exception("Error adding a new taxon to the taxonomy")
        else :
            return taxid
        
79

80
    cpdef close(self) :
81
        if (obi_close_taxonomy(self._pointer) < 0) :
82
            raise Exception("Error closing the taxonomy")
83 84
    
    
85

86
cdef class OBI_Taxon :    # TODO dict subclass?
87 88

    def __init__(self, object taxon_capsule) :
89 90 91 92 93 94 95 96
        self._pointer = <ecotx_t*> PyCapsule_GetPointer(taxon_capsule, NULL)
        if self._pointer == NULL :
            raise Exception("Error reading the taxonomy")

    # name property getter
    @property
    def name(self):
        return bytes2str(self._pointer.name)
97
    
98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115
    # taxid property getter
    @property
    def taxid(self):
        return self._pointer.taxid

    # rank property getter
    @property
    def rank(self):
        return self._pointer.rank
   
    # farest property getter
    @property
    def farest(self):
        return self._pointer.farest

    # parent property getter
    @property
    def parent(self):
116
        cdef object parent_capsule
117
        parent_capsule = PyCapsule_New(self._pointer.parent, NULL, NULL)
118
        return OBI_Taxon(parent_capsule)
119 120 121 122 123 124 125 126

    def __repr__(self):
        d = {}
        d['taxid']   = self.taxid
        d['name']    = self.name
        d['parent']  = self.parent.taxid
        d['farest']  = self.farest
        return str(d)
127 128
    
    
129