Commit 9dfebbd2 authored by Martin Larralde's avatar Martin Larralde
Browse files

Implement pickle protocol for `pyrodigal.Genes`

parent a3befd28
......@@ -272,6 +272,8 @@ cdef class Genes:
cdef readonly TrainingInfo training_info
cpdef size_t __sizeof__(self)
cpdef dict __getstate__(self)
cpdef object __setstate__(self, dict state)
cdef inline _gene* _add_gene(
self,
......
......@@ -181,6 +181,8 @@ class Genes(typing.Sequence[Gene]):
def __iter__(self) -> Iterator[Gene]: ...
def __reversed__(self) -> Iterator[Gene]: ...
def __sizeof__(self) -> int: ...
def __getstate__(self) -> Dict[str, object]: ...
def __setstate__(self, state: Dict[str, object]) -> None: ...
def clear(self) -> None: ...
def write_gff(self, file: TextIO, prefix: str = "gene_") -> int: ...
def write_genes(self, file: TextIO, prefix: str ="gene_", width: typing.Optional[int] = 70) -> int: ...
......
......@@ -2960,6 +2960,52 @@ cdef class Genes:
cpdef size_t __sizeof__(self):
return self.capacity * sizeof(_gene) + sizeof(self)
cpdef dict __getstate__(self):
cdef size_t i
return {
"_num_seq": self._num_seq,
"nodes": self.nodes,
"sequence": self.sequence,
"training_info": self.training_info,
"genes": [
{
"begin": self.genes[i].begin,
"end": self.genes[i].end,
"start_ndx": self.genes[i].start_ndx,
"stop_ndx": self.genes[i].stop_ndx,
}
for i in range(self.length)
]
}
cpdef object __setstate__(self, dict state):
cdef size_t i
cdef dict gene
cdef list genes = state["genes"]
# realloc to the exact number of genes
self.length = self.capacity = len(genes)
if self.capacity > 0:
self.genes = <_gene*> PyMem_Realloc(self.genes, self.capacity * sizeof(_gene))
if self.genes == NULL:
raise MemoryError("Failed to reallocate gene array")
else:
PyMem_Free(self.genes)
self.genes = NULL
# copy attributes
self._num_seq = state["_num_seq"]
self.nodes = state["nodes"]
self.sequence = state["sequence"]
self.training_info = state["training_info"]
# copy gene data from the state dictionary
for i, gene in enumerate(genes):
self.genes[i].begin = gene["begin"]
self.genes[i].end = gene["end"]
self.genes[i].start_ndx = gene["start_ndx"]
self.genes[i].stop_ndx = gene["stop_ndx"]
# --- C interface --------------------------------------------------------
cdef inline _gene* _add_gene(
......
......@@ -6,6 +6,7 @@ import os
import platform
import sys
import unittest
import pickle
from .. import OrfFinder
from .fasta import parse
......@@ -94,3 +95,8 @@ class TestGenes(unittest.TestCase):
# can occur between different platforms: for instance,
# depending on the compilation flags, the score
# `-32.2550` may be rounded as `-32.25` or `-32.26`.
def test_pickle(self):
genes = pickle.loads(pickle.dumps(self.genes))
for gene1, gene2 in zip(self.genes, genes):
self.assertEqual(gene1._gene_data, gene2._gene_data)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment