Skip to content

Commit

Permalink
Deserialize appose_type "shm" and "ndarray"
Browse files Browse the repository at this point in the history
  • Loading branch information
tpietzsch committed Apr 22, 2024
1 parent c17d7c6 commit 55f3e23
Showing 1 changed file with 32 additions and 1 deletion.
33 changes: 32 additions & 1 deletion src/appose/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@

import json
from typing import Any, Dict
from multiprocessing import shared_memory

Args = Dict[str, Any]

Expand All @@ -38,4 +39,34 @@ def encode(data: Args) -> str:


def decode(the_json: str) -> Args:
return json.loads(the_json)
return json.loads(the_json, object_hook=_appose_object_hook)


class ShmNDArray:

def __init__(self, shm: shared_memory.SharedMemory, dtype: str, shape ):
self.shm = shm
self.dtype = dtype
self.shape = shape

def __str__(self):
return f"ShmNDArray(shm='{self.shm.name}' ({self.shm.size}), dtype='{self.dtype}', shape={self.shape})"

def ndarray(self):
try:
import math
import numpy
num_elements = math.prod(self.shape)
return numpy.ndarray(num_elements, dtype=self.dtype, buffer=self.shm.buf).reshape(self.shape)
except ModuleNotFoundError:
raise ImportError("NumPy is not available.")


def _appose_object_hook(obj: Dict):
type = obj.get('appose_type')
if type == 'shm':
return shared_memory.SharedMemory(name=(obj['name']), size=(obj['size']))
elif type == 'ndarray':
return ShmNDArray(obj['shm'], obj['dtype'], obj['shape'])
else:
return obj

0 comments on commit 55f3e23

Please sign in to comment.