Fix genome mapping serialization

Actually test and fix issues with using the new serialization based
approach for dealing with blat hits.

Author: blakesweeney

rnacentral_pipeline/cli/genome_mapping.py

@@ -37,41 +37,42 @@ def hits():
     pass
 
 
-@hits.command('as-json')
+@hits.command('serialize')
 @click.argument('assembly_id')
 @click.argument('hits', default='-', type=click.File('r'))
-@click.argument('output', default='-', type=click.File('w'))
+@click.argument('output', default='-', type=click.File('wb'))
 def hits_json(assembly_id, hits, output):
     """
-    Convert the PSL file into a JSON-line file (one object per line). This is a
-    lossy operation but keeps everything needed for selecting later.
+    Serialize the PSL file into something that python can later process. This is
+    a lossy operation but keeps everything needed for selecting later. This
+    exists so we can do mulitple select steps and still merge the results.
     """
-    blat.as_json(assembly_id, hits, output)
+    blat.as_pickle(assembly_id, hits, output)
 
 
-@cli.command('as-importable')
-@click.argument('hits', default='-', type=click.File('r'))
+@hits.command('as-importable')
+@click.argument('hits', default='-', type=click.File('rb'))
 @click.argument('output', default='-', type=click.File('w'))
-def as_importable(raw, output):
+def as_importable(hits, output):
     """
     Convert a json-line file into a CSV that can be used for import by pgloader.
     This is lossy as it only keeps the things needed for the database.
     """
-    blat.as_importable(raw, output)
+    blat.write_importable(hits, output)
 
 
 @hits.command('select')
-@click.option('--sort', default=False)
-@click.argument('hits', default='-', type=click.File('r'))
-@click.argument('output', default='-', type=click.File('w'))
+@click.option('--sort', is_flag=True, default=False)
+@click.argument('hits', default='-', type=click.File('rb'))
+@click.argument('output', default='-', type=click.File('wb'))
 def select_hits(hits, output, sort=False):
     """
     Parse a JSON-line file and select the best hits in the file. The best hits
     are written to the output file. This assumes the file is sorted by
     urs_taxid unless --sort is given in which case the data is sorted in memory.
     That may be very expensive.
     """
-    blat.select_json(hits, output, sort=sort)
+    blat.select_pickle(hits, output, sort=sort)
 
 
 @cli.command('url-for')

rnacentral_pipeline/databases/data/regions.py

@@ -74,7 +74,7 @@ def display_int(self):
 
 
 # @enum.unique
-class CoordianteStart(enum.Enum):
+class CoordinateStart(enum.Enum):
     zero = 0
     one = 1
 
@@ -124,9 +124,19 @@ class CoordinateSystem(object):
     http://genome.ucsc.edu/blog/the-ucsc-genome-browser-coordinate-counting-systems/
     """
 
-    basis = attr.ib(validator=is_a(CoordianteStart))
+    basis = attr.ib(validator=is_a(CoordinateStart))
     close_status = attr.ib(validator=is_a(CloseStatus))
 
+    @classmethod
+    def build(cls, value):
+        if isinstance(value, six.text_type):
+            return cls.from_name(value)
+        if isinstance(value, dict):
+            return cls(**value)
+        if isinstance(value, cls):
+            return value
+        raise ValueError("Cannot build CoordinateSystem from %s" % str(value))
+
     @classmethod
     def from_name(cls, name):
         """
@@ -143,7 +153,7 @@ def from_name(cls, name):
             raise UnknownCoordinateSystem(name)
 
         return cls(
-            basis=CoordianteStart.from_name(basis_name),
+            basis=CoordinateStart.from_name(basis_name),
             close_status=CloseStatus.from_name(close_name),
         )
 
@@ -178,9 +188,9 @@ def size(self, location):
 
     def as_zero_based(self, location):
         start = location.start
-        if self.basis is CoordianteStart.zero:
+        if self.basis is CoordinateStart.zero:
             pass
-        elif self.basis is CoordianteStart.one:
+        elif self.basis is CoordinateStart.one:
             start = start - 1
         else:
             raise ValueError("Unknown type of start: %s" % self.basis)
@@ -189,9 +199,9 @@ def as_zero_based(self, location):
 
     def as_one_based(self, location):
         start = location.start
-        if self.basis is CoordianteStart.zero:
+        if self.basis is CoordinateStart.zero:
             start = start + 1
-        elif self.basis is CoordianteStart.one:
+        elif self.basis is CoordinateStart.one:
             pass
         else:
             raise ValueError("Unknown type of start: %s" % self.basis)
@@ -218,16 +228,26 @@ def greater_than_start(self, attribute, value):
                              (value, self.start))
 
 
+def as_sorted_exons(raw):
+    exons = []
+    for entry in raw:
+        if isinstance(entry, dict):
+           exons.append(Exon(**entry))
+        else:
+            exons.append(entry)
+    return tuple(sorted(exons, key=op.attrgetter('start')))
+
+
 @attr.s(frozen=True, hash=True, slots=True)
 class SequenceRegion(object):
     assembly_id = attr.ib(validator=is_a(six.text_type), converter=six.text_type)
     chromosome = attr.ib(validator=is_a(six.text_type), converter=six.text_type)
     strand = attr.ib(validator=is_a(Strand), converter=Strand.build)
-    exons = attr.ib(
-        validator=is_a(tuple), 
-        converter=lambda es: tuple(sorted(es, key=op.attrgetter('start'))),
+    exons = attr.ib(validator=is_a(tuple), converter=as_sorted_exons)
+    coordinate_system = attr.ib(
+        validator=is_a(CoordinateSystem),
+        converter=CoordinateSystem.build,
     )
-    coordinate_system = attr.ib(validator=is_a(CoordinateSystem))
 
     @property
     def start(self):

rnacentral_pipeline/rnacentral/genome_mapping/blat.py

@@ -14,6 +14,7 @@
 """
 
 import csv
+import json
 import operator as op
 import itertools as it
 import logging
@@ -23,6 +24,7 @@
 import attr
 from attr.validators import instance_of as is_a
 
+from rnacentral_pipeline import utils
 from rnacentral_pipeline.databases.data.regions import Exon
 from rnacentral_pipeline.databases.data.regions import Strand
 from rnacentral_pipeline.databases.data.regions import SequenceRegion
@@ -129,12 +131,6 @@ def parse_psl(assembly_id, handle):
         yield BlatHit.build(assembly_id, result)
 
 
-def parse_json(handle):
-    for line in handle:
-        data = json.loads(line)
-        yield BlatHit(**data)
-
-
 def select_hits(hits):
     hits = six.moves.filter(select_possible, hits)
     hits = it.groupby(hits, op.attrgetter('upi'))
@@ -144,27 +140,21 @@ def select_hits(hits):
     return hits
 
 
-def write_json(hits, output):
-    for hit in hits:
-        output.write(json.dumps(hit))
-        output.write('\n')
-
-
-def write_importable(hits, output):
+def write_importable(handle, output):
+    hits = utils.unpickle_stream(handle)
     writeable = six.moves.map(op.methodcaller('writeable'), hits)
-    writeable = it.chain.from_iterable(hits)
+    writeable = it.chain.from_iterable(writeable)
     csv.writer(output).writerows(writeable)
 
 
-def as_json(assembly_id, hits, output):
+def as_pickle(assembly_id, hits, output):
     parsed = parse_psl(assembly_id, hits)
-    parsed = six.moves.map(attr.asdict, selected)
-    write_json(parsed, output)
+    utils.pickle_stream(parsed, output)
 
 
-def select_json(hits, output, sort=False):
-    parsed = parse_json(hits)
+def select_pickle(handle, output, sort=False):
+    hits = utils.unpickle_stream(handle)
     if sort:
-        parsed = sorted(parsed, key=op.itemgetter('upi'))
-    selected = select_hits(parsed)
-    write_json(selected, output)
+        hits = sorted(hits, key=op.attrgetter('upi'))
+    selected = select_hits(hits)
+    utils.pickle_stream(selected, output)

rnacentral_pipeline/utils.py

@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+
+"""
+Copyright [2009-2018] EMBL-European Bioinformatics Institute
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import six.moves.cPickle as pickle
+
+def pickle_stream(stream, handle, *args, **kwargs):
+    for entry in stream:
+        pickle.dump(entry, handle, *args, **kwargs)
+
+
+def unpickle_stream(handle, *args, **kwargs):
+    try:
+        while True:
+            yield pickle.load(handle)
+    except EOFError:
+        raise StopIteration()

tests/databases/data/regions_test.py

@@ -13,6 +13,8 @@
 limitations under the License.
 """
 
+import six.moves.cPickle as pickle
+
 import pytest
 
 from rnacentral_pipeline.databases.data.regions import *
@@ -46,12 +48,12 @@ def test_fails_with_bad_strands(raw):
 @pytest.mark.parametrize('name,expected', [
     ('0-start, half-open', 
      CoordinateSystem(
-         basis=CoordianteStart.zero, 
+         basis=CoordinateStart.zero, 
          close_status=CloseStatus.open)
      ),
     ('1-start, fully-closed', 
      CoordinateSystem(
-         basis=CoordianteStart.one,
+         basis=CoordinateStart.one,
          close_status=CloseStatus.closed)
      ),
 ])
@@ -234,3 +236,40 @@ def test_can_generate_correct_writeable(accession, strand, coordinate_system, ex
     )
     upi = accession.startswith('U')
     assert list(region.writeable(accession, is_upi=upi)) == expected
+
+
+@pytest.mark.parametrize('data', [
+    (Strand.reverse),
+    (Strand.unknown),
+    (Strand.forward),
+    (CoordinateStart.zero),
+    (CoordinateStart.one),
+    (CloseStatus.closed),
+    (CloseStatus.open),
+])
+def test_can_serialize_enums_to_and_from_pickle(data):
+    assert pickle.loads(pickle.dumps(data)) is data
+
+
+@pytest.mark.parametrize('region', [
+    (SequenceRegion(
+        assembly_id='GRCh38',
+        chromosome='1',
+        strand=Strand.unknown,
+        exons=[Exon(start=100000, stop=2000000), Exon(start=-1, stop=10000)],
+        coordinate_system=CoordinateSystem.from_name('0-start, half-open')
+    )),
+    (SequenceRegion(
+        assembly_id='GRCh38',
+        chromosome='MT',
+        strand=Strand.forward,
+        exons=[Exon(start=10, stop=200), Exon(start=-1, stop=10000)],
+        coordinate_system=CoordinateSystem.from_name('1-start, fully-closed')
+    )),
+])
+def test_can_serialize_regions_to_and_from_pickle(region):
+    data = attr.asdict(region)
+    loaded = pickle.loads(pickle.dumps(data)) 
+    assert loaded == data
+    assert SequenceRegion(**loaded) == region
+    assert pickle.loads(pickle.dumps(region)) == region

tests/utils_test.py

@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+
+"""
+Copyright [2009-2018] EMBL-European Bioinformatics Institute
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import tempfile
+
+import six
+
+from rnacentral_pipeline import utils
+
+def test_can_serialize_stream():
+    data = ['a', 1, 2, 3, 4]
+    with tempfile.NamedTemporaryFile() as tmp:
+        utils.pickle_stream(data, tmp)
+        tmp.seek(0)
+        result = utils.unpickle_stream(tmp)
+        assert not isinstance(result, list)
+        for index, obj in enumerate(result):
+            assert data[index] == obj