rawfile-localpv/util.py

70 lines
1.5 KiB
Python
Raw Normal View History

2020-04-24 11:48:09 +00:00
import base64
2020-04-23 13:52:01 +00:00
import functools
2020-04-24 11:48:09 +00:00
import inspect
import pickle
2020-04-23 13:52:01 +00:00
import subprocess
def indent(obj, lvl):
return "\n".join([(lvl * " ") + line for line in str(obj).splitlines()])
def log_grpc_request(func):
@functools.wraps(func)
def wrap(self, request, context):
try:
res = func(self, request, context)
print(
f"""{func.__name__}({{
2020-04-23 13:52:01 +00:00
{indent(request, 2)}
}}) = {{
{indent(res, 2)}
}}"""
)
return res
except Exception as exc:
ret = (str(context._state.code), context._state.details)
print(
f"""{func.__name__}({{
{indent(request, 2)}
}}) = {ret}
"""
)
raise exc
2020-04-23 13:52:01 +00:00
return wrap
def run(cmd):
return subprocess.run(cmd, shell=True, check=True)
2020-04-24 11:48:09 +00:00
def run_out(cmd: str):
p = subprocess.run(cmd, shell=True, capture_output=True)
return p
2020-04-24 11:48:09 +00:00
class remote_fn(object):
def __init__(self, fn):
self.fn = fn
def as_cmd(self, *args, **kwargs):
call_data = [inspect.getsource(self.fn).encode(), args, kwargs]
call_data_serialized = base64.b64encode(pickle.dumps(call_data))
run_cmd = f"""
python <<EOF
import base64
import pickle
remote_fn = lambda fn: fn # FIXME: dirty hack
call_data = pickle.loads(base64.b64decode({call_data_serialized}))
exec(call_data[0])
{self.fn.__name__}(*call_data[1], **call_data[2])
EOF
"""
return run_cmd
def __call__(self, *args, **kwargs):
raise Exception("Should only be run inside pod")