/anypytools/h5py_wrapper.py

https://github.com/AnyBody-Research-Group/AnyPyTools
Python | 164 lines | 126 code | 27 blank | 11 comment | 37 complexity | 679b64308b3d548e0cacbecaf3ef32b1 MD5 | raw file
  1. # -*- coding: utf-8 -*-
  2. """
  3. Created on Mon Jan 16 11:40:42 2012.
  4. @author: mel
  5. """
  6. import logging
  7. import h5py
  8. logger = logging.getLogger("abt.anypytools")
  9. def _follow_reftarget(elem):
  10. completename = elem.attrs["CompleteName"]
  11. completename = completename.replace(b".", b"/")
  12. reftarget = elem.attrs["RefTarget"].replace(b".", b"/")
  13. prefix = completename[: -len(elem.name)]
  14. h5target = reftarget[len(prefix) :]
  15. elem = elem.file[h5target]
  16. return elem
  17. def _check_input_path(path):
  18. """Convert dot notation to stardard h5py path."""
  19. if "/" not in path:
  20. # path does not have traditional h5 format.
  21. if path.startswith("Main.") and "Output" in path:
  22. path = "/Output" + path.split("Output")[-1]
  23. path = path.replace(".", "/")
  24. return path
  25. class File(h5py.File): # noqa
  26. __doc__ = h5py.File.__doc__
  27. def __init__(self, *args, **kwargs): # noqa
  28. super(File, self).__init__(*args, **kwargs)
  29. self.wrapped = True
  30. def __getitem__(self, path): # noqa
  31. """."""
  32. path = _check_input_path(path)
  33. try:
  34. elem = super(File, self).file[path]
  35. if isinstance(elem, h5py.Group) and not len(elem.keys()):
  36. if "RefTarget" in elem.attrs:
  37. elem = _follow_reftarget(elem)
  38. except KeyError:
  39. elem = super(type(self), self)
  40. levels = path.strip("/").split("/")
  41. for level in levels:
  42. if elem.__contains__(level):
  43. elem = elem.__getitem__(level)
  44. else:
  45. try:
  46. elem = _follow_reftarget(elem)
  47. elem = elem.__getitem__(level)
  48. except Exception:
  49. raise KeyError("Entry not found: " + path)
  50. if isinstance(elem, h5py.Group):
  51. return Group(elem.id)
  52. elif isinstance(elem, h5py.Dataset):
  53. return Dataset(elem.id)
  54. elif isinstance(elem, h5py.File):
  55. return File(elem.id)
  56. @property
  57. def file(self): # noqa
  58. id = super(File, self).file.id
  59. return File(id)
  60. @property
  61. def parent(self): # noqa
  62. id = super(File, self).parent.id
  63. return Group(id)
  64. def __contains__(self, name): # noqa
  65. """Test if a member name exists"""
  66. if super(File, self).__contains__(name):
  67. return True
  68. else:
  69. try:
  70. self.__getitem__(name)
  71. return True
  72. except KeyError:
  73. pass
  74. return False
  75. class Group(h5py.Group): # noqa
  76. __doc__ = h5py.Group.__doc__
  77. def __init__(self, arg): # noqa
  78. super(Group, self).__init__(arg)
  79. self.wrapped = True
  80. def __getitem__(self, path): # noqa
  81. path = _check_input_path(path)
  82. try:
  83. elem = super(Group, self).__getitem__(path)
  84. if isinstance(elem, h5py.Group) and not len(elem.keys()):
  85. if "RefTarget" in elem.attrs:
  86. elem = _follow_reftarget(elem)
  87. except KeyError:
  88. elem = super(type(self), self)
  89. levels = path.strip("/").split("/")
  90. for level in levels:
  91. if elem.__contains__(level):
  92. elem = elem.__getitem__(level)
  93. else:
  94. try:
  95. elem = _follow_reftarget(elem)
  96. elem = elem.__getitem__(level)
  97. except Exception:
  98. raise KeyError("Entry not found: " + path)
  99. if isinstance(elem, h5py.Group):
  100. return Group(elem.id)
  101. elif isinstance(elem, h5py.Dataset):
  102. return Dataset(elem.id)
  103. elif isinstance(elem, h5py.File):
  104. return File(elem.id)
  105. @property
  106. def file(self): # noqa
  107. id = super(Group, self).file.id
  108. return File(id)
  109. @property
  110. def parent(self): # noqa
  111. id = super(Group, self).parent.id
  112. return Group(id)
  113. def __contains__(self, name): # noqa
  114. """Test if a member name exists"""
  115. if super(Group, self).__contains__(name):
  116. return True
  117. else:
  118. try:
  119. self.__getitem__(name)
  120. return True
  121. except KeyError:
  122. pass
  123. return False
  124. class Dataset(h5py.Dataset): # noqa
  125. __doc__ = h5py.Dataset.__doc__
  126. def __init__(self, arg): # noqa
  127. super(Dataset, self).__init__(arg)
  128. self.wrapped = True
  129. @property
  130. def file(self): # noqa
  131. id = super(Dataset, self).file.id
  132. return File(id)
  133. @property
  134. def parent(self): # noqa
  135. id = super(Dataset, self).parent.id
  136. return Group(id)