pyinfra/test/fixtures/input.py

40 lines
874 B
Python

from itertools import repeat
import numpy as np
import pytest
from PIL import Image
from funcy import lmap, compose
from test.utils.image import image_to_bytes
@pytest.fixture
def data(data_type, pdf):
if data_type == "pdf":
return pdf
elif data_type == "bytestring":
return "content".encode("latin1")
@pytest.fixture
def data_items(item_type):
if item_type == "string":
return [bytes(f"content{i}", encoding="utf8") for i in range(7)]
elif item_type == "image":
return images()
else:
raise ValueError(f"Unknown item type {item_type}")
def array_to_image(array) -> Image.Image:
return Image.fromarray(np.uint8(array * 255), mode="RGB")
def input_batch():
return np.random.random_sample(size=(7, 3, 30, 30))
def images():
return lmap(compose(image_to_bytes, array_to_image), input_batch())