Skip to content

Commit

Permalink
feat: Python File interface types support
Browse files Browse the repository at this point in the history
  • Loading branch information
thewtex committed Oct 20, 2022
1 parent 8cc00d9 commit b9509ba
Show file tree
Hide file tree
Showing 6 changed files with 77 additions and 11 deletions.
4 changes: 2 additions & 2 deletions src/python/itkwasm/itkwasm/binary_file.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from dataclasses import dataclass
from pathlib import PurePosixPath

@dataclass
class BinaryFile:
data: bytes
path: str
path: PurePosixPath
28 changes: 25 additions & 3 deletions src/python/itkwasm/itkwasm/pipeline.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
import json
from pathlib import Path
from typing import List, Union, Dict
from typing import List, Union, Dict, Tuple
from .interface_types import InterfaceTypes
from .pipeline_input import PipelineInput
from .pipeline_output import PipelineOutput
from .text_stream import TextStream
from .binary_stream import BinaryStream
from .text_file import TextFile
from .binary_file import BinaryFile

from wasmer import engine, wasi, Store, Module, ImportObject, Instance
from wasmer_compiler_cranelift import Compiler
Expand All @@ -26,9 +28,21 @@ def __init__(self, pipeline: Union[str, Path, bytes]):
self.module = Module(self.store, wasm_bytes)
self.wasi_version = wasi.get_version(self.module, strict=True)

def run(self, args: List[str], outputs: List[PipelineOutput]=[], inputs: List[PipelineInput]=[], preopen_directories=[], map_directories={}, environments={}) -> List[PipelineOutput]:
def run(self, args: List[str], outputs: List[PipelineOutput]=[], inputs: List[PipelineInput]=[]) -> Tuple[PipelineOutput]:
"""Run the itk-wasm pipeline."""

preopen_directories=set()
map_directories={}
# Todo: expose?
environments={}
for index, input_ in enumerate(inputs):
if input_.type == InterfaceTypes.TextFile or input_.type == InterfaceTypes.BinaryFile:
preopen_directories.add(str(input_.data.path.parent))
for index, output in enumerate(outputs):
if output.type == InterfaceTypes.TextFile or output.type == InterfaceTypes.BinaryFile:
preopen_directories.add(str(output.data.path.parent))
preopen_directories = list(preopen_directories)

wasi_state = wasi.StateBuilder('itk-wasm-pipeline')
wasi_state.arguments(args)
wasi_state.environments(environments)
Expand Down Expand Up @@ -60,6 +74,10 @@ def run(self, args: List[str], outputs: List[PipelineOutput]=[], inputs: List[Pi
array_ptr = self._set_input_array(data_array, index, 0)
data_json = { "size": len(data_array), "data": f"data:application/vnd.itk.address,0:{array_ptr}" }
self._set_input_json(data_json, index)
elif input_.type == InterfaceTypes.TextFile:
pass
elif input_.type == InterfaceTypes.BinaryFile:
pass
else:
raise ValueError(f'Unexpected/not yet supported input.type {input_.type}')

Expand All @@ -80,13 +98,17 @@ def run(self, args: List[str], outputs: List[PipelineOutput]=[], inputs: List[Pi
data_size = self.output_array_size(0, index, 0)
data_array = bytes(memoryview(self.memory.buffer)[data_ptr:data_ptr+data_size])
output_data = PipelineOutput(InterfaceTypes.BinaryStream, BinaryStream(data_array))
elif output.type == InterfaceTypes.TextFile:
output_data = PipelineOutput(InterfaceTypes.TextFile, TextFile(output.data.path))
elif output.type == InterfaceTypes.BinaryFile:
output_data = PipelineOutput(InterfaceTypes.BinaryFile, BinaryFile(output.data.path))
populated_outputs.append(output_data)

delayed_exit = instance.exports.itk_wasm_delayed_exit
delayed_exit(return_code)

# Should we be returning the return_code?
return populated_outputs
return tuple(populated_outputs)

def _set_input_array(self, data_array: bytes, input_index: int, sub_index: int) -> int:
data_ptr = 0
Expand Down
4 changes: 2 additions & 2 deletions src/python/itkwasm/itkwasm/text_file.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from dataclasses import dataclass
from pathlib import PurePosixPath

@dataclass
class TextFile:
data: str
path: str
path: PurePosixPath
1 change: 1 addition & 0 deletions src/python/itkwasm/test/input/input.bin
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
ޭ��
1 change: 1 addition & 0 deletions src/python/itkwasm/test/input/input.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
The answer is 42.
50 changes: 46 additions & 4 deletions src/python/itkwasm/test/test_pipeline.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
from pathlib import Path
from pathlib import Path, PurePosixPath

from itkwasm import InterfaceTypes, TextStream, BinaryStream, PipelineInput, PipelineOutput, Pipeline
from itkwasm import InterfaceTypes, TextStream, BinaryStream, PipelineInput, PipelineOutput, Pipeline, TextFile, BinaryFile

test_input_dir = Path(__file__).resolve().parent / 'input'
import tempfile


def test_stdout_stderr():
Expand All @@ -16,7 +17,7 @@ def test_pipeline_bytes():
pipeline = Pipeline(wasm_bytes)
pipeline.run([])

def test_pipeline_input_output_files():
def test_pipeline_input_output_streams():
pipeline = Pipeline(test_input_dir / 'input-output-files-test.wasi.wasm')

pipeline_inputs = [
Expand Down Expand Up @@ -45,4 +46,45 @@ def test_pipeline_input_output_files():
assert outputs[1].data.data[0], 222
assert outputs[1].data.data[1], 173
assert outputs[1].data.data[2], 190
assert outputs[1].data.data[3], 239
assert outputs[1].data.data[3], 239

def test_pipeline_input_output_files():
pipeline = Pipeline(test_input_dir / 'input-output-files-test.wasi.wasm')
input_text_file = PurePosixPath(test_input_dir / 'input.txt')
input_binary_file = PurePosixPath(test_input_dir / 'input.bin')

with tempfile.TemporaryDirectory() as tmpdirname:
output_text_file = PurePosixPath(Path(tmpdirname) / 'output.txt')
output_binary_file = PurePosixPath(Path(tmpdirname) / 'output.bin')

pipeline_inputs = [
PipelineInput(InterfaceTypes.TextFile, TextFile(input_text_file)),
PipelineInput(InterfaceTypes.BinaryFile, BinaryFile(input_binary_file)),
]

pipeline_outputs = [
PipelineOutput(InterfaceTypes.TextFile, TextFile(output_text_file)),
PipelineOutput(InterfaceTypes.BinaryFile, BinaryFile(output_binary_file)),
]

args = [
'--use-files',
'--input-text-file', str(input_text_file),
'--input-binary-file', str(input_binary_file),
'--output-text-file', str(output_text_file),
'--output-binary-file', str(output_binary_file)
]

outputs = pipeline.run(args, pipeline_outputs, pipeline_inputs)

assert outputs[0].type == InterfaceTypes.TextFile
with open(outputs[0].data.path, 'r') as fp:
content = fp.read()
assert content == 'The answer is 42.'
assert outputs[1].type, InterfaceTypes.BinaryFile
with open(outputs[1].data.path, 'rb') as fp:
content = fp.read()
assert content[0] == 222
assert content[1] == 173
assert content[2] == 190
assert content[3] == 239

0 comments on commit b9509ba

Please sign in to comment.