Skip to content

Commit

Permalink
Change zip to fail if iteratables are different lengths
Browse files Browse the repository at this point in the history
  • Loading branch information
tomwhite committed Nov 4, 2024
1 parent 528eea0 commit 0e18357
Show file tree
Hide file tree
Showing 3 changed files with 9 additions and 9 deletions.
6 changes: 3 additions & 3 deletions tests/test_vcf_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ def test_write_vcf__filtering(tmp_path, include, exclude, expected_chrom_pos):
assert len(variants) == len(expected_chrom_pos)
assert v.samples == ["NA00001", "NA00002", "NA00003"]

for variant, chrom_pos in zip(variants, expected_chrom_pos, strict=False):
for variant, chrom_pos in zip(variants, expected_chrom_pos, strict=True):
chrom, pos = chrom_pos
assert variant.CHROM == chrom
assert variant.POS == pos
Expand Down Expand Up @@ -141,7 +141,7 @@ def test_write_vcf__regions(tmp_path, regions, targets,

assert v.samples == ["NA00001", "NA00002", "NA00003"]

for variant, chrom_pos in zip(variants, expected_chrom_pos, strict=False):
for variant, chrom_pos in zip(variants, expected_chrom_pos, strict=True):
chrom, pos = chrom_pos
assert variant.CHROM == chrom
assert variant.POS == pos
Expand Down Expand Up @@ -230,7 +230,7 @@ def test_write_vcf__regions_samples_filtering(
assert len(variants) == len(expected_chrom_pos)
assert v.samples == ["NA00001"]

for variant, chrom_pos in zip(variants, expected_chrom_pos, strict=False):
for variant, chrom_pos in zip(variants, expected_chrom_pos, strict=True):
chrom, pos = chrom_pos
assert variant.CHROM == chrom
assert variant.POS == pos
Expand Down
10 changes: 5 additions & 5 deletions vcztools/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ def generate(root):
end = start + v_chunk_size

for gt_row, phase in zip(
gt_zarray[start:end], phase_zarray[start:end], strict=False
gt_zarray[start:end], phase_zarray[start:end], strict=True
):

def stringify(gt_and_phase: tuple):
Expand All @@ -113,7 +113,7 @@ def stringify(gt_and_phase: tuple):
return separator.join(gt)

gt_row = gt_row.tolist()
yield map(stringify, zip(gt_row, phase, strict=False))
yield map(stringify, zip(gt_row, phase, strict=True))
else:
# TODO: Support datasets without the phasing data
raise NotImplementedError
Expand Down Expand Up @@ -219,8 +219,8 @@ def _compose_sample_loop_generator(

def generate(root):
iterables = (generator(root) for generator in generators)
zipped = zip(*iterables, strict=False)
zipped_zipped = (zip(*element, strict=False) for element in zipped)
zipped = zip(*iterables, strict=True)
zipped_zipped = (zip(*element, strict=True) for element in zipped)
flattened_zipped_zipped = (
(
subsubelement
Expand Down Expand Up @@ -309,7 +309,7 @@ def generate(root) -> str:
filter_iterable = filter_generator(root)

for results, filter_indicator in zip(
zip(*iterables, strict=False), filter_iterable, strict=False
zip(*iterables, strict=True), filter_iterable, strict=True
):
if filter_indicator:
results = map(str, results)
Expand Down
2 changes: 1 addition & 1 deletion vcztools/stats.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ def stats(vcz, output):
).astype(np.int64)

for contig, contig_length, nr in zip(
contigs, contig_lengths, num_records_per_contig, strict=False
contigs, contig_lengths, num_records_per_contig, strict=True
):
if nr > 0:
print(f"{contig}\t{contig_length}\t{nr}", file=output)

0 comments on commit 0e18357

Please sign in to comment.