134 lines
3.7 KiB
Python
134 lines
3.7 KiB
Python
import tempfile
|
|
from PIL import Image
|
|
import numpy
|
|
from pathlib import Path
|
|
import pytest
|
|
|
|
from pymol import cmd
|
|
|
|
|
|
def datafile(filename: str) -> Path:
|
|
"""
|
|
Return path to filename, the current directory and the data
|
|
directory are searched for filename.
|
|
:filename: filename to search for
|
|
"""
|
|
if Path(filename).exists():
|
|
return filename
|
|
PYMOL_TESTING_ROOT = 2
|
|
# pymol-testing root
|
|
pymol_test_dir = Path(__file__).parents[PYMOL_TESTING_ROOT]
|
|
return pymol_test_dir.joinpath('data', filename)
|
|
|
|
|
|
class mktemp(object):
|
|
"""
|
|
Context manager which returns a temporary filename and
|
|
deletes the file in the end, if it exists.
|
|
"""
|
|
|
|
def __init__(self, suffix: str = ''):
|
|
self.file = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
|
|
self.file.close() # Allow other processes to access the file
|
|
self.filename = self.file.name
|
|
|
|
def __enter__(self):
|
|
return self.filename
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
if Path(self.filename).exists():
|
|
Path(self.filename).unlink()
|
|
|
|
|
|
def ambientOnly(cmd) -> None:
|
|
"""
|
|
Set up a scene with only ambient lighting.
|
|
"""
|
|
cmd.set('ambient', 1)
|
|
cmd.set('antialias', 0)
|
|
cmd.set('light_count', 1)
|
|
cmd.set('depth_cue', 0)
|
|
# needed for open-source
|
|
cmd.set('reflect', 0)
|
|
cmd.set('direct', 0)
|
|
|
|
|
|
def get_imagearray(cmd, **kwargs) -> numpy.array:
|
|
"""
|
|
Return the image as a numpy array.
|
|
:kwargs: keyword arguments passed to cmd.png
|
|
"""
|
|
|
|
with mktemp('.png') as filename:
|
|
cmd.png(filename, **kwargs)
|
|
img = Image.open(filename)
|
|
|
|
return numpy.array(img.getdata(),
|
|
numpy.uint8).reshape((img.size[1], img.size[0], -1))
|
|
|
|
|
|
def _imageHasColor(cmd, color, img, delta=0) -> bool:
|
|
"""
|
|
Return True if the image contains the given color.
|
|
:color: color to search for
|
|
:img: image as numpy array
|
|
:delta: maximum difference between color and image color
|
|
"""
|
|
if isinstance(color, str):
|
|
color = [int(v*255) for v in cmd.get_color_tuple(color)]
|
|
else:
|
|
color = list(color)
|
|
dim = img.shape[-1]
|
|
if dim == len(color) + 1:
|
|
dim -= 1
|
|
img = img[..., :dim]
|
|
diff = abs(img.reshape((-1, dim)) - color)
|
|
return (diff - delta <= 0).prod(1).sum()
|
|
|
|
|
|
def imageHasColor(cmd, color: str, delta: float = 0) -> bool:
|
|
"""
|
|
Return True if the image contains the given color.
|
|
:color: color to search for
|
|
:delta: maximum difference between color and image color
|
|
"""
|
|
img = get_imagearray(cmd)
|
|
return _imageHasColor(cmd, color, img, delta)
|
|
|
|
|
|
def assert_in_names_undo(cmd, name: str) -> None:
|
|
"""
|
|
Assert that the object is in the names list and undoing the
|
|
command removes it from the list.
|
|
:name: name to search for
|
|
"""
|
|
assert name in cmd.get_names()
|
|
cmd.undo()
|
|
assert name not in cmd.get_names()
|
|
cmd.redo()
|
|
assert name in cmd.get_names()
|
|
|
|
|
|
def compatible_with(version: str) -> bool:
|
|
def tupleize_version(str_: str):
|
|
return tuple(int(x) for x in str_.split('.') if x.isdigit())
|
|
|
|
PYMOL_VERSION = cmd.get_version()
|
|
PYMOL_VERSION_TUPLE = tupleize_version(PYMOL_VERSION[0])
|
|
|
|
if isinstance(version, int):
|
|
return version <= PYMOL_VERSION[2]
|
|
elif isinstance(version, float):
|
|
return version <= PYMOL_VERSION[1]
|
|
else:
|
|
return tupleize_version(version) <= PYMOL_VERSION_TUPLE
|
|
|
|
|
|
def requires_version(version, reason=None):
|
|
def decorator(test_func):
|
|
return pytest.mark.skipif(
|
|
not compatible_with(version),
|
|
reason=reason or f"Requires PyMOL {version}"
|
|
)(test_func)
|
|
return decorator
|