import torch
import io
import json
from petrel_client.client import Client

client = Client('./conf/petreloss.conf')

# Pytorch save & load

data = torch.tensor([0, 1, 2, 3])
tensor_url = 's3://bucket1/tensor_data'

with io.BytesIO() as f:
    torch.save(data, f)
    client.put(tensor_url, f.getvalue())

with io.BytesIO(client.get(tensor_url)) as f:
    data2 = torch.load(f)

assert torch.equal(data, data2)


# Json dumps & loads

data = [0, 1, 2, 3]
json_data_url = 's3://bucket1/json_data'

client.put(json_data_url, json.dumps(data).encode('utf-8'))

data2 = json.loads(client.get(json_data_url))

assert data == data2