added context manager for reading / writing #9

.pre-commit-config.yaml

@@ -0,0 +1,41 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.4.0
+    hooks:
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: check-yaml
+      - id: check-added-large-files
+      - id: check-merge-conflict
+      - id: debug-statements
+      - id: name-tests-test
+        args: ["--pytest-test-first"]
+      - id: no-commit-to-branch
+        args: ["--branch", "main", "--branch", "develop"]
+  - repo: https://github.com/psf/black
+    rev: 23.3.0
+    hooks:
+      - id: black
+  - repo: https://github.com/pycqa/isort
+    rev: 5.12.0
+    hooks:
+      - id: isort
+        name: isort (python)
+        args: ["--profile", "black", "--filter-files"]
+  - repo: https://github.com/asottile/pyupgrade
+    rev: v3.4.0
+    hooks:
+      - id: pyupgrade
+        args: ["--py38-plus"]
+  - repo: https://github.com/PyCQA/autoflake
+    rev: v2.1.1
+    hooks:
+      - id: autoflake
+        args: [
+          "--in-place",
+          "--remove-all-unused-imports",
+          "--ignore-init-module-imports",
+          "--expand-star-imports",
+        ]

heparchy/__init__.py

@@ -4,11 +4,56 @@
 
 The hierarchically formatted high energy physics IO library.
 """
+import contextlib as ctx
+import typing as ty
 import warnings
+from pathlib import Path
 
 from heparchy import read, write
 from heparchy._version import __version__, __version_tuple__
 
-
-__all__ = ["read", "write", "__version__", "__version_tuple__"]
+__all__ = ["read", "write", "open_file", "__version__", "__version_tuple__"]
 warnings.filterwarnings("once", category=DeprecationWarning)
+
+
+@ty.overload
+def open_file(
+    path: ty.Union[str, Path], mode: ty.Literal["r"]
+) -> ty.ContextManager[read.hdf.HdfReader]:
+    ...
+
+
+@ty.overload
+def open_file(
+    path: ty.Union[str, Path], mode: ty.Literal["w"]
+) -> ty.ContextManager[write.hdf.HdfWriter]:
+    ...
+
+
+@ctx.contextmanager
+def open_file(path: ty.Union[str, Path], mode: ty.Literal["r", "w"]):
+    """High level file manager for reading and writing HEP data to HDF5.
+
+    Parameters
+    ----------
+    path : Path or str
+        Path on disk where the file is to be opened.
+    mode : {'r', 'w'}
+        Whether to open the file in 'r' (read) or 'w' (write) mode.
+
+    Yields
+    ------
+    HdfReader or HdfWriter
+        Heparchy's file handler for accessing the HDF5 file.
+    """
+    if mode not in {"r", "w"}:
+        raise ValueError(f'Mode {mode} not known. Please use "r" or "w".')
+    stack = ctx.ExitStack()
+    if mode == "r":
+        f = stack.enter_context(read.hdf.HdfReader(path))
+    elif mode == "w":
+        f = stack.enter_context(write.hdf.HdfWriter(path))
+    try:
+        yield f
+    finally:
+        stack.close()

heparchy/read/base.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from typing import List, Any, Dict
+from typing import Any
 
 import numpy as np
 import numpy.typing as npt
 
-
-DoubleVector = npt.NDArray[np.float64]
 IntVector = npt.NDArray[np.int32]
 HalfIntVector = npt.NDArray[np.int16]
+DoubleVector = npt.NDArray[np.float64]
 BoolVector = npt.NDArray[np.bool_]
+AnyVector = npt.NDArray[Any]
+VoidVector = npt.NDArray[np.void]
 
 
 class EventReaderBase(ABC):
@@ -50,7 +52,7 @@ def final(self) -> BoolVector:
 
     @property
     @abstractmethod
-    def available(self) -> List[str]:
+    def available(self) -> list[str]:
         """Provides list of all dataset names in event."""
 
     @abstractmethod
@@ -78,7 +80,7 @@ def string(self) -> str:
 
     @property
     @abstractmethod
-    def decay(self) -> Dict[str, IntVector]:
+    def decay(self) -> dict[str, IntVector]:
         """Returns dictionary with two entries, describing the hard
         interaction for this process.