def make_read_writer(outfile, contigs=None): """Creates a writer to outfile writing Read protos. This function creates an writer that accepts Read protos and writes them to outfile. The type of the writer is determined by the extension of outfile. If it's one of SAM_EXTENSIONS, we will write out SAM records via make_sam_writer. Otherwise we will write out TFRecord file of serialized Read protos. Args: outfile: A path to a file where we want to write our variant calls. contigs: A list of the reference genome contigs for writers that need contig information. Returns: An writer object and a write_fn accepting a Read proto that writes to writer. Raises: ValueError: If any of the optional arguments needed for the specific output type of outfile are missing. """ if any(outfile.endswith(ext) for ext in SAM_EXTENSIONS): if contigs is None: raise ValueError('contigs must be provided for SAM/BAM output') raise NotImplementedError else: return io_utils.RawProtoWriterAdaptor( io_utils.make_tfrecord_writer(outfile))
def test_adaptor_with_ownership(self, take_ownership): mock_writer = mock.MagicMock() adaptor = io.RawProtoWriterAdaptor(mock_writer, take_ownership=take_ownership) # Write out protos to our adaptor. with adaptor as enter_return_value: # Make sure that __enter__ returns the adaptor itself. self.assertIs(adaptor, enter_return_value) adaptor.write(self.proto1) adaptor.write(self.proto2) if take_ownership: # If we took ownership, mock_writer __enter__ and __exit__ should have # been called. mock_writer.__enter__.assert_called_once_with() test_utils.assert_called_once_workaround(mock_writer.__exit__) else: # If not, they shouldn't have been called. test_utils.assert_not_called_workaround(mock_writer.__enter__) test_utils.assert_not_called_workaround(mock_writer.__exit__) self.assertEqual( mock_writer.write.call_args_list, [mock.call(r.SerializeToString()) for r in self.protos])
def __init__(self, options): self._writers = {k: None for k in ['candidates', 'examples', 'gvcfs']} if options.candidates_filename: self._add_writer( 'candidates', io_utils.RawProtoWriterAdaptor( io_utils.make_tfrecord_writer( options.candidates_filename))) if options.examples_filename: self._add_writer( 'examples', io_utils.RawProtoWriterAdaptor( io_utils.make_tfrecord_writer(options.examples_filename))) if options.gvcf_filename: self._add_writer( 'gvcfs', io_utils.RawProtoWriterAdaptor( io_utils.make_tfrecord_writer(options.gvcf_filename)))
def make_variant_writer(outfile, contigs, samples=None, filters=None): """Creates a writer to outfile writing variant Protos. This function creates an writer that accepts Variant protos and writes them to outfile. The type of the writer is determined by the extension of outfile. If it's one of VCF_EXTENSIONS, we will write out VCF records via make_vcf_writer. Otherwise we will write out TFRecord file of serialized Variant protos. Args: outfile: A path to a file where we want to write our variant calls. contigs: A list of the reference genome contigs for writers that need contig information. samples: A list of sample names we will be writing out. Can be None if the list of samples isn't required for the intended output type. Will raise an exception is None and is required. filters: A list of VcfFilterInfo protos defining the filter fields present in the to-be-written records. Can be None if the intended writer doesn't require this information, but will raise an exception if None and required by the destination writer. Returns: An writer object that supports the context manager protocol for opening/closing the writer and a single function write(proto). Raises: ValueError: If any of the optional arguments needed for the specific output type of outfile are missing. """ if any(outfile.endswith(ext) for ext in VCF_EXTENSIONS): if samples is None: raise ValueError('samples must be provided for vcf output') if filters is None: raise ValueError('filters must be provided for vcf output') return make_vcf_writer( outfile, contigs=contigs, samples=samples, filters=filters) else: return io_utils.RawProtoWriterAdaptor( io_utils.make_tfrecord_writer(outfile))