From 4cbc457780ce1f25789818fb5a83c10a4b0fd279 Mon Sep 17 00:00:00 2001 From: dachengx Date: Sat, 16 Sep 2023 12:10:33 -0500 Subject: [PATCH] Mandatorily wrap _read_chunk in a check_chunk_n decorator --- strax/storage/common.py | 7 +++++++ strax/storage/files.py | 1 - strax/storage/mongo.py | 1 - strax/storage/zipfiles.py | 2 -- 4 files changed, 7 insertions(+), 4 deletions(-) diff --git a/strax/storage/common.py b/strax/storage/common.py index fb8a1a50a..514e7ec3c 100644 --- a/strax/storage/common.py +++ b/strax/storage/common.py @@ -403,6 +403,13 @@ class StorageBackend: these have to be hardcoded (or made part of the key). """ + def __new__(cls, *args, **kwargs): + """Mandatorily wrap _read_chunk in a check_chunk_n decorator""" + if '_read_chunk' in cls.__dict__: + method = getattr(cls, '_read_chunk') + setattr(cls, '_read_chunk', strax.check_chunk_n(method)) + return super(StorageBackend, cls).__new__(cls) + def loader(self, backend_key, time_range=None, diff --git a/strax/storage/files.py b/strax/storage/files.py index 7db4e7e44..12e5d06a8 100644 --- a/strax/storage/files.py +++ b/strax/storage/files.py @@ -251,7 +251,6 @@ def _read_and_format_chunk(self, *args, **kwargs): chunk.target_size_mb = self.set_chunk_size_mb return chunk - @strax.check_chunk_n def _read_chunk(self, dirname, chunk_info, dtype, compressor): fn = osp.join(dirname, chunk_info['filename']) return strax.load_file(fn, dtype=dtype, compressor=compressor) diff --git a/strax/storage/mongo.py b/strax/storage/mongo.py index 71ede79c6..621427541 100644 --- a/strax/storage/mongo.py +++ b/strax/storage/mongo.py @@ -45,7 +45,6 @@ def __init__(self, uri: str, database: str, col_name: str): self._buff_mb = DEFAULT_MONGO_BACKEND_BUFFER_MB self._buff_nruns = DEFAULT_MONGO_BACKEND_BUFFER_NRUNS - @strax.check_chunk_n def _read_chunk(self, backend_key, chunk_info, dtype, compressor): """See strax.Backend""" chunk_i = chunk_info["chunk_i"] diff --git a/strax/storage/zipfiles.py b/strax/storage/zipfiles.py index 688117fc7..c4beb9b4b 100644 --- a/strax/storage/zipfiles.py +++ b/strax/storage/zipfiles.py @@ -104,8 +104,6 @@ def zip_dir(input_dir, output_zipfile, delete=False): @export class ZipFileBackend(strax.StorageBackend): - - @strax.check_chunk_n def _read_chunk(self, zipn_and_dirn, chunk_info, dtype, compressor): zipn, dirn = zipn_and_dirn with zipfile.ZipFile(zipn) as zp: