diff --git a/heparchy/__init__.py b/heparchy/__init__.py index abb2b31..5895ba0 100644 --- a/heparchy/__init__.py +++ b/heparchy/__init__.py @@ -23,6 +23,20 @@ def open_file( ... +@ty.overload +def open_file( + path: ty.Union[str, Path], mode: ty.Literal["r"], process: None +) -> ty.ContextManager[read.hdf.HdfReader]: + ... + + +@ty.overload +def open_file( + path: ty.Union[str, Path], mode: ty.Literal["r"], process: str +) -> ty.ContextManager[read.hdf.HdfProcessReader]: + ... + + @ty.overload def open_file( path: ty.Union[str, Path], mode: ty.Literal["w"] @@ -30,8 +44,26 @@ def open_file( ... +@ty.overload +def open_file( + path: ty.Union[str, Path], mode: ty.Literal["w"], process: None +) -> ty.ContextManager[write.hdf.HdfWriter]: + ... + + +@ty.overload +def open_file( + path: ty.Union[str, Path], mode: ty.Literal["w"], process: str +) -> ty.ContextManager[write.hdf.HdfProcessWriter]: + ... + + @ctx.contextmanager -def open_file(path: ty.Union[str, Path], mode: ty.Literal["r", "w"]): +def open_file( + path: ty.Union[str, Path], + mode: ty.Literal["r", "w"], + process: ty.Optional[str] = None, +): """High level file manager for reading and writing HEP data to HDF5. Parameters @@ -40,10 +72,13 @@ def open_file(path: ty.Union[str, Path], mode: ty.Literal["r", "w"]): 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. + process : str, optional + If provided, will provide the handler to the process indexed + with the value. Yields ------ - HdfReader or HdfWriter + HdfReader or HdfProcessReader or HdfWriter or HdfProcessWriter Heparchy's file handler for accessing the HDF5 file. """ if mode not in {"r", "w"}: @@ -51,9 +86,17 @@ def open_file(path: ty.Union[str, Path], mode: ty.Literal["r", "w"]): stack = ctx.ExitStack() if mode == "r": f = stack.enter_context(read.hdf.HdfReader(path)) + if process is not None: + f = f[process] elif mode == "w": f = stack.enter_context(write.hdf.HdfWriter(path)) + if process is not None: + f = stack.enter_context(f.new_process(process)) try: yield f finally: stack.close() + + +with open_file("/scratch/jlc1n20/graftnet/data/h_bb.hdf5", "w", "test") as f: + f