1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
|
"""
Read and write on compressed files.
"""
import os
import os.path
import numpy as np
import pytest
from ase import io
from ase.io import formats
from ase.build import bulk
single = bulk('Au')
multiple = [bulk('Fe'), bulk('Zn'), bulk('Li')]
compressions = ['gz', 'bz2', 'xz']
def test_get_compression():
"""Identification of supported compression from filename."""
assert formats.get_compression('H2O.pdb.gz') == ('H2O.pdb', 'gz')
assert formats.get_compression('CH4.pdb.bz2') == ('CH4.pdb', 'bz2')
assert formats.get_compression('Alanine.pdb.xz') == ('Alanine.pdb', 'xz')
# zip not implemented ;)
assert formats.get_compression('DNA.pdb.zip') == ('DNA.pdb.zip', None)
assert formats.get_compression('crystal.cif') == ('crystal.cif', None)
@pytest.mark.parametrize('ext', compressions)
def test_compression_write_single(ext):
"""Writing compressed file."""
filename = 'single.xsf.{ext}'.format(ext=ext)
io.write(filename, single)
assert os.path.exists(filename)
@pytest.mark.parametrize('ext', compressions)
def test_compression_read_write_single(ext):
"""Re-reading a compressed file."""
# Use xsf filetype as it needs to check the 'magic'
# filetype guessing when reading
filename = 'single.xsf.{ext}'.format(ext=ext)
io.write(filename, single)
assert os.path.exists(filename)
reread = io.read(filename)
assert reread.get_chemical_symbols() == single.get_chemical_symbols()
assert np.allclose(reread.positions, single.positions)
@pytest.mark.parametrize('ext', compressions)
def test_compression_write_multiple(ext):
"""Writing compressed file, with multiple configurations."""
filename = 'multiple.xyz.{ext}'.format(ext=ext)
io.write(filename, multiple)
assert os.path.exists(filename)
@pytest.mark.parametrize('ext', compressions)
def test_compression_read_write_multiple(ext):
"""Re-reading a compressed file with multiple configurations."""
filename = 'multiple.xyz.{ext}'.format(ext=ext)
io.write(filename, multiple)
assert os.path.exists(filename)
reread = io.read(filename, ':')
assert len(reread) == len(multiple)
assert np.allclose(reread[-1].positions, multiple[-1].positions)
@pytest.mark.parametrize('ext', compressions)
def test_modes(ext):
"""Test the different read/write modes for a compression format."""
filename = 'testrw.{ext}'.format(ext=ext)
for mode in ['w', 'wb', 'wt']:
with formats.open_with_compression(filename, mode) as tmp:
if 'b' in mode:
tmp.write(b'some text')
else:
tmp.write('some text')
for mode in ['r', 'rb', 'rt']:
with formats.open_with_compression(filename, mode) as tmp:
if 'b' in mode:
assert tmp.read() == b'some text'
else:
assert tmp.read() == 'some text'
|
