Added ability to limit size of file cache in LH5Store

src/lgdo/lh5/iterator.py

@@ -77,7 +77,7 @@ def __init__(
             The friend should have the same length and entry list. A single
             LH5 table containing columns from both iterators will be returned.
         """
-        self.lh5_st = LH5Store(base_path=base_path, keep_open=True)
+        self.lh5_st = LH5Store(base_path=base_path, keep_open=3)
 
         # List of files, with wildcards and env vars expanded
         if isinstance(lh5_files, str):
@@ -182,9 +182,7 @@ def _get_file_cumlen(self, i_file: int) -> int:
             fcl = self.file_map[i_start - 1] if i_start > 0 else 0
 
             for i in range(i_start, i_file + 1):
-                fcl += self.lh5_st.read_n_rows(
-                    self.groups[i], self.lh5_files[i]
-                )
+                fcl += self.lh5_st.read_n_rows(self.groups[i], self.lh5_files[i])
                 self.file_map[i] = fcl
         return fcl
 

src/lgdo/lh5/store.py

@@ -8,6 +8,7 @@
 import logging
 import os
 import sys
+from collections import OrderedDict
 from collections.abc import Mapping, Sequence
 from inspect import signature
 from typing import Any
@@ -46,14 +47,15 @@ def __init__(
             directory path to prepend to LH5 files.
         keep_open
             whether to keep files open by storing the :mod:`h5py` objects as
-            class attributes.
+            class attributes. If ``keep_open`` is an ``int``, keep only the
+            ``n`` most recently opened files; if ``True``, no limit
         locking
             whether to lock files when reading
         """
         self.base_path = "" if base_path == "" else utils.expand_path(base_path)
         self.keep_open = keep_open
         self.locking = locking
-        self.files = {}
+        self.files = OrderedDict()
 
     def gimme_file(
         self,
@@ -86,6 +88,7 @@ def gimme_file(
             file_kwargs["locking"] = self.locking
 
         if lh5_file in self.files:
+            self.files.move_to_end(lh5_file)
             return self.files[lh5_file]
 
         if self.base_path != "":
@@ -119,6 +122,8 @@ def gimme_file(
         h5f = h5py.File(full_path, mode, **file_kwargs)
 
         if self.keep_open:
+            if isinstance(self.keep_open, int) and len(self.files) >= self.keep_open:
+                self.files.popitem(last=False)
             self.files[lh5_file] = h5f
 
         return h5f