jsonify instead of pickle

gaurav879 · Jun 11, 2021 · af28e62 · af28e62
1 parent f86e74d
commit af28e62
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 22 deletions.
diff --git a/conftest.py b/conftest.py
@@ -189,16 +189,13 @@ def _storage_from_request(request, memory_storage, local_storage, s3_storage):
     cache_sizes = []
 
     if MEMORY in requested_providers:
-        _skip_if_none(memory_storage)
         storage_providers.append(memory_storage)
         cache_sizes.append(MIN_FIRST_CACHE_SIZE)
     if LOCAL in requested_providers:
-        _skip_if_none(local_storage)
         storage_providers.append(local_storage)
         cache_size = MIN_FIRST_CACHE_SIZE if not cache_sizes else MIN_SECOND_CACHE_SIZE
         cache_sizes.append(cache_size)
     if S3 in requested_providers:
-        _skip_if_none(s3_storage)
         storage_providers.append(s3_storage)
 
     if len(storage_providers) == len(cache_sizes):
@@ -211,18 +208,21 @@ def _storage_from_request(request, memory_storage, local_storage, s3_storage):
 def memory_storage(request):
     if not _is_opt_true(request, MEMORY_OPT):
         return _get_memory_provider(request)
+    pytest.skip()
 
 
 @pytest.fixture
 def local_storage(request):
     if _is_opt_true(request, LOCAL_OPT):
         return _get_local_provider(request)
+    pytest.skip()
 
 
 @pytest.fixture
 def s3_storage(request):
     if _is_opt_true(request, S3_OPT):
         return _get_s3_provider(request)
+    pytest.skip()
 
 
 @pytest.fixture
@@ -232,19 +232,16 @@ def storage(request, memory_storage, local_storage, s3_storage):
 
 @pytest.fixture
 def memory_ds(memory_storage):
-    _skip_if_none(memory_storage)
     return _get_dataset(memory_storage)
 
 
 @pytest.fixture
 def local_ds(local_storage):
-    _skip_if_none(local_storage)
     return _get_dataset(local_storage)
 
 
 @pytest.fixture
 def s3_ds(s3_storage):
-    _skip_if_none(s3_storage)
     return _get_dataset(s3_storage)
 
 
@@ -310,4 +307,4 @@ def skip_if_full_benchmarks_disabled(request, marks):
         return
 
     if FULL_BENCHMARK_MARK in marks:
-        pytest.skip()
+        pytest.skip()
diff --git a/hub/core/meta/dataset_meta.py b/hub/core/meta/dataset_meta.py
@@ -1,12 +1,12 @@
-import pickle  # TODO: NEVER USE PICKLE
+import json
 
 from hub.core.typing import StorageProvider
 from hub.util.keys import get_dataset_meta_key
 
 
 def write_dataset_meta(storage: StorageProvider, meta: dict):
-    storage[get_dataset_meta_key()] = pickle.dumps(meta)
+    storage[get_dataset_meta_key()] = json.dumps(meta)
 
 
 def read_dataset_meta(storage: StorageProvider) -> dict:
-    return pickle.loads(storage[get_dataset_meta_key()])
+    return json.loads(storage[get_dataset_meta_key()])
diff --git a/hub/core/meta/index_map.py b/hub/core/meta/index_map.py
@@ -1,4 +1,4 @@
-import pickle  # TODO: NEVER USE PICKLE
+import json
 from typing import List
 
 from hub.core.typing import StorageProvider
@@ -7,8 +7,8 @@
 
 def write_index_map(key: str, storage: StorageProvider, index_map: list):
     index_map_key = get_index_map_key(key)
-    storage[index_map_key] = pickle.dumps(index_map)
+    storage[index_map_key] = json.dumps(index_map)
 
 
 def read_index_map(key: str, storage: StorageProvider) -> List[dict]:
-    return pickle.loads(storage[get_index_map_key(key)])
+    return json.loads(storage[get_index_map_key(key)])
diff --git a/hub/core/meta/tensor_meta.py b/hub/core/meta/tensor_meta.py
@@ -1,4 +1,4 @@
-import pickle  # TODO: NEVER USE PICKLE
+import json
 from typing import Any, Callable, Optional
 
 import numpy as np
@@ -11,11 +11,11 @@
 
 
 def write_tensor_meta(key: str, storage: StorageProvider, meta: dict):
-    storage[get_tensor_meta_key(key)] = pickle.dumps(meta)
+    storage[get_tensor_meta_key(key)] = json.dumps(meta)
 
 
 def read_tensor_meta(key: str, storage: StorageProvider) -> dict:
-    return pickle.loads(storage[get_tensor_meta_key(key)])
+    return json.loads(storage[get_tensor_meta_key(key)])
 
 
 def default_tensor_meta(
@@ -48,8 +48,8 @@ def update_tensor_meta_with_array(
     if batched:
         shape = shape[1:]
     tensor_meta["dtype"] = str(array.dtype)
-    tensor_meta["min_shape"] = shape
-    tensor_meta["max_shape"] = shape
+    tensor_meta["min_shape"] = list(shape)
+    tensor_meta["max_shape"] = list(shape)
 
     return tensor_meta
 

diff --git a/hub/core/tensor.py b/hub/core/tensor.py
@@ -201,5 +201,6 @@ def _check_array_and_tensor_are_compatible(tensor_meta: dict, array: np.ndarray)
 
 
 def _update_tensor_meta_shapes(shape: Tuple[int], tensor_meta: dict):
-    tensor_meta["min_shape"] = tuple(np.minimum(tensor_meta["min_shape"], shape))
-    tensor_meta["max_shape"] = tuple(np.maximum(tensor_meta["max_shape"], shape))
+    for i, dim in enumerate(shape):
+        tensor_meta["min_shape"][i] = min(dim, tensor_meta["min_shape"][i])
+        tensor_meta["max_shape"][i] = max(dim, tensor_meta["max_shape"][i])
diff --git a/hub/core/tests/common.py b/hub/core/tests/common.py
@@ -177,8 +177,8 @@ def run_engine_test(
                 "chunk_size": chunk_size,
                 "length": sample_count,
                 "dtype": a_in.dtype.name,
-                "min_shape": tuple(expected_min_shape),
-                "max_shape": tuple(expected_max_shape),
+                "min_shape": list(expected_min_shape),
+                "max_shape": list(expected_max_shape),
             },
         )