Coverage for src/metador_core/container/drivers.py: 100%

1"""Metador driver abstraction in order to enable different underlying implementations.""" 

2from enum import Enum 

3from typing import Any, Optional, Type, Union, cast 

4 

5import h5py 

6 

7from ..ih5.container import IH5Record 

8from ..util.types import H5FileLike, OpenMode 

9 

10 

11class MetadorDriverEnum(Enum): 

12 """Supported classes that work with MetadorContainer. 

13 

14 Note that they must be unrelated (i.e. not subclasses of each other). 

15 """ 

16 

17 HDF5 = h5py.File 

18 IH5 = IH5Record 

19 

20 @classmethod 

21 def to_dict(cls): 

22 return {x: x.value for x in iter(cls)} 

23 

24 

25# NOTE: must be duplicated or static checkers can't pick it up 

26MetadorDriver = Union[h5py.File, IH5Record] 

27"""Union of all supported classes (for static type check).""" 

28 

29METADOR_DRIVERS = MetadorDriverEnum.to_dict() 

30"""Dict representation of MetadorDriverEnum.""" 

31 

32METADOR_DRIVER_CLASSES = tuple(METADOR_DRIVERS.values()) 

33"""Tuple of all supported classes (for instance check).""" 

34 

35# ---- 

36 

37 

38def get_driver_type(raw_cont: MetadorDriver) -> MetadorDriverEnum: 

39 """Return the driver type of container (if it is a suitable (sub)class).""" 

40 for val, cls in METADOR_DRIVERS.items(): 

41 if isinstance(raw_cont, cls): 

42 return val 

43 raise ValueError(f"Object not of known container driver type: {raw_cont}") 

44 

45 

46def get_source(raw_cont: MetadorDriver, driver: MetadorDriverEnum = None) -> Any: 

47 """Return an object (i.e. input resource(s)) needed to re-open the given container.""" 

48 c = cast(Any, raw_cont) 

49 driver = driver or get_driver_type(raw_cont) 

50 if driver == MetadorDriverEnum.HDF5: 

51 return c.filename 

52 elif driver == MetadorDriverEnum.IH5: 

53 return c.ih5_files 

54 

55 

56def to_h5filelike( 

57 name_or_obj: Union[MetadorDriver, Any], 

58 mode: OpenMode = "r", 

59 *, 

60 # NOTE: driver takes actual class instead of enum, to also allow subclasses 

61 driver: Optional[Type[MetadorDriver]] = None, 

62) -> H5FileLike: 

63 """Given a container or a resource with a driver, try to return a H5FileLike. 

64 

65 If first argument is instance of a known driver is returned unchanged. 

66 Otherwise, will try to open it using the driver (h5py.File by default). 

67 

68 Returns a H5FileLike compatible with MetadorContainer, or raises ValueError. 

69 """ 

70 if isinstance(name_or_obj, METADOR_DRIVER_CLASSES): 

71 # user has passed a h5file-like object already 

72 return cast(H5FileLike, name_or_obj) 

73 else: 

74 # user passed arguments to try instantiating a container object 

75 driver = driver or h5py.File 

76 if not issubclass(driver, METADOR_DRIVER_CLASSES): 

77 msg = f"Passed driver class not supported: {driver}" 

78 raise ValueError(msg) 

79 return cast(H5FileLike, driver(cast(Any, name_or_obj), mode))