Simplify pickling

pyensembl/ensembl_release.py

@@ -90,12 +90,9 @@ def __hash__(self):
         return hash((self.release, self.species))
 
     def __getstate__(self):
-        fields = Genome.__getstate__(self)
-        fields["release"] = self.release
-        fields["species"] = self.species
-        return fields
+        # Must be in order of __init__ arguments
+        return [self.release, self.species, self.server]
 
     def __setstate__(self, fields):
-        # Genome sets __dict__ equal to all fields, so the release and species
-        # fields are handled as a part of that.
-        Genome.__setstate__(self, fields)
+        self.__init__(*fields)
+

pyensembl/gene.py

@@ -46,6 +46,7 @@ def __init__(
 
         Locus.__init__(self, contig, start, end, strand)
 
+        self.require_valid_biotype = require_valid_biotype
         if require_valid_biotype and not is_valid_biotype(biotype):
             raise ValueError(
                 "Invalid gene_biotype %s for gene with ID = %s" % (
@@ -74,14 +75,12 @@ def __hash__(self):
         return hash(self.id)
 
     def __getstate__(self):
-        fields = self.__dict__.copy()
-        # We can't pickle connections
-        del fields["db"]
-        return fields
+        # Must be in order of __init__ arguments
+        return [self.id, self.name, self.contig, self.start, self.end, self.strand,
+                self.biotype, self.genome, self.require_valid_biotype]
 
     def __setstate__(self, fields):
-        self.__dict__ = fields
-        self.db = self.genome.db
+        self.__init__(*fields)
 
     @memoized_property
     def transcripts(self):

pyensembl/genome.py

@@ -98,9 +98,6 @@ def __init__(
         self._transcript_fasta_path_or_url = transcript_fasta_path_or_url
         self._protein_fasta_path_or_url = protein_fasta_path_or_url
 
-        self._init()
-
-    def _init(self):
         self.download_cache = DownloadCache(
             reference_name=self.reference_name,
             annotation_name=self.annotation_name,
@@ -121,6 +118,17 @@ def _init(self):
 
         self._init_lazy_fields()
 
+    def __getstate__(self):
+        # Must be in order of __init__ arguments
+        return [self.reference_name, self.annotation_name, self.annotation_version,
+                self._gtf_path_or_url, self._transcript_fasta_path_or_url,
+                self._protein_fasta_path_or_url, self.decompress_on_download,
+                self.copy_local_files_to_cache, self.require_ensembl_ids,
+                self.cache_directory_path]
+
+    def __setstate__(self, fields):
+        self.__init__(*fields)
+
     def _init_lazy_fields(self):
         """
         Member data that gets loaded or constructed on demand
@@ -1042,25 +1050,3 @@ def protein_ids(self, contig=None, strand=None):
             distinct=True)
         # drop None values
         return [protein_id for protein_id in protein_ids if protein_id]
-
-    def __getstate__(self):
-        # Not the same as _fields(); these are useful for pickling/unpickling even if not necessary
-        # when checking for Genome equality.
-        field_list = [
-            "reference_name",
-            "annotation_name",
-            "annotation_version",
-            "_gtf_path_or_url",
-            "_transcript_fasta_path_or_url",
-            "_protein_fasta_path_or_url",
-            "decompress_on_download",
-            "copy_local_files_to_cache",
-            "require_ensembl_ids",
-            "cache_directory_path"]
-        fields = self.__dict__.copy()
-        fields = dict([(field, value) for (field, value) in fields.items() if field in field_list])
-        return fields
-
-    def __setstate__(self, fields):
-        self.__dict__ = fields
-        self._init()

pyensembl/transcript.py

@@ -41,6 +41,7 @@ def __init__(
             gene_id,
             genome,
             require_valid_biotype=True):
+        self.require_valid_biotype = require_valid_biotype
         if require_valid_biotype and not is_valid_biotype(biotype):
             raise ValueError(
                 "Invalid biotype '%s' for transcript with ID=%s, name=%s" % (
@@ -94,14 +95,12 @@ def gene(self):
         return self.genome.gene_by_id(self.gene_id)
 
     def __getstate__(self):
-        fields = self.__dict__.copy()
-        # We can't pickle connections
-        del fields["db"]
-        return fields
+        # Must be in order of __init__ arguments
+        return [self.id, self.name, self.contig, self.start, self.end, self.strand,
+                self.biotype, self.gene_id, self.genome, self.require_valid_biotype]
 
     def __setstate__(self, fields):
-        self.__dict__ = fields
-        self.db = self.genome.db
+        self.__init__(*fields)
 
     @memoized_property
     def exons(self):