Bug fix

legend-exp · Oct 31, 2023 · 052e175 · 052e175
1 parent eef36c3
commit 052e175
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 12 deletions.
diff --git a/src/lgdo/__init__.py b/src/lgdo/__init__.py
@@ -41,7 +41,9 @@
 from ._version import version as __version__
 
 # from .lh5 import Iterator, Store, load_dfs, load_nda, ls, show
-from .lh5_store import LH5Iterator, LH5Store, load_dfs, load_nda, ls, show
+from .lh5 import Iterator, Store, load_dfs, load_nda, ls, show
+
+# from .lh5_store import LH5Iterator, LH5Store
 from .types import (
     LGDO,
     Array,
@@ -68,11 +70,13 @@
     "VectorOfVectors",
     "VectorOfEncodedVectors",
     "WaveformTable",
-    "LH5Iterator",
-    "LH5Store",
+    "Iterator",
+    "Store",
     "load_dfs",
     "load_nda",
     "ls",
     "show",
+    #   "LH5Iterator",
+    #   "LH5Store"
     "__version__",
 ]
diff --git a/src/lgdo/lh5_store.py b/src/lgdo/lh5_store.py
@@ -19,19 +19,18 @@
 from .types import VectorOfVectors  # noqa: F401
 from .types import WaveformTable  # noqa: F401
 
-warn(
-    "lgdo.lh5_store has moved to a subfolder lgdo.lh5 containing Store and Iterator. "
-    "Please replace 'import lgdo.lh5_store' with 'import lgdo.lh5'. "
-    "lgdo.lh5_store will be removed in a future release.",
-    DeprecationWarning,
-    stacklevel=2,
-)
-
-
 DEFAULT_HDF5_COMPRESSION = None
 
 
 class LH5Iterator(lh5.Iterator):
+    warn(
+        "lgdo.lh5_store has moved to a subfolder lgdo.lh5 containing Store and Iterator. "
+        "Please replace 'from lgdo.lh5_store import LH5Iterator' with 'from lgdo.lh5 import Iterator'. "
+        "lgdo.lh5_store will be removed in a future release.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+
     def __init__(
         self,
         lh5_files: str | list[str],
@@ -56,6 +55,14 @@ def __init__(
 
 
 class LH5Store(lh5.Store):
+    warn(
+        "lgdo.lh5_store has moved to a subfolder lgdo.lh5 containing Store and Iterator. "
+        "Please replace 'from lgdo.lh5_store import LH5Store' with 'from lgdo.lh5 import Store'. "
+        "lgdo.lh5_store will be removed in a future release.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+
     def __init__(self, base_path: str = "", keep_open: bool = False):
         super().__init__(base_path, keep_open)