Skip to content

Commit

Permalink
Add: Type-checking safeguards
Browse files Browse the repository at this point in the history
  • Loading branch information
ashvardanian committed May 19, 2023
1 parent dd02645 commit 1b866ae
Showing 1 changed file with 10 additions and 7 deletions.
17 changes: 10 additions & 7 deletions src/ucall/_server.py
Original file line number Diff line number Diff line change
@@ -1,21 +1,24 @@
import numpy as np
from PIL import Image
from typing import Callable, get_type_hints
import inspect
from typing import Callable, get_type_hints
from functools import wraps
from io import BytesIO

import numpy as np
from PIL import Image


class _Server:
server = None

def __init__(self) -> None:
self.server = None

def __call__(self, func: Callable):
return self.route(func)

def run(self, max_cycles: int = -1, max_seconds: float = -1):
return self.server.run(max_cycles, max_seconds)

def unpack(self, arg: bytes, hint):
def unpack(self, arg: bytes, hint: type):
if hint == bytes or hint == bytearray:
return arg

Expand All @@ -36,8 +39,7 @@ def pack(self, res):
buf = BytesIO()
if not res.format:
res.format = 'tiff'
res.save(buf, res.format, compression='raw',
compression_level=0)
res.save(buf, res.format, compression='raw', compression_level=0)

return buf.getvalue()

Expand All @@ -52,6 +54,7 @@ def wrapper(*args, **kwargs):
new_kwargs = {}

for arg, hint in zip(args, hints.values()):
assert isinstance(hint, type), 'Hint must be a type!'
if isinstance(arg, bytes):
new_args.append(self.unpack(arg, hint))
else:
Expand Down

0 comments on commit 1b866ae

Please sign in to comment.