CondHDF5ESSource/scripts/condhdf5tohdf5.py

0001 #!/usr/bin/env python3
0002 import argparse
0003 import sys
0004 import logging
0005 import copy
0006 import h5py
0007 import numpy as np
0008 from collections import OrderedDict
0009 import zlib
0010 import lzma
0011 from CondCore.CondHDF5ESSource.hdf5Writer import writeH5File
0012
0013 #Global tags hold a list of Tags
0014 # Tags give the
0015 #      record name,
0016 #      list of data products
0017 #      list of IOVs
0018 #      list of payloads per IOV
0019 # Payloads give
0020 #      a payload name and
0021 #      the serialized data for a data product
0022 #      the type of data for the data product
0023 #
0024
0025
0026 class IOVSyncValue(object):
0027     def __init__(self, high, low):
0028         self.high = high
0029         self.low = low
0030
0031 class H5Payload(object):
0032     def __init__(self,dataset,name, compressor):
0033         self._dataset = dataset
0034         self._hash = name
0035         self._type = dataset.attrs['type']
0036         self._memsize = dataset.attrs['memsize']
0037         self._compressor = compressor
0038     def name(self):
0039         return self._hash
0040     def actualType(self):
0041         return self._type
0042     def memsize(self):
0043         return self._memsize
0044     def data(self):
0045         ds = self._dataset[()]
0046         if len(ds) == self.memsize():
0047             return ds
0048         #was compressed
0049         return self._compressor.decompress(ds)
0050
0051 class H5DataProduct(object):
0052     def __init__(self, group, name, compressor):
0053         self._type = group.attrs['type']
0054         self._name = name
0055         self._payloadGroup = group['Payloads']
0056         self._compressor = compressor
0057     def name(self):
0058         return self._name
0059     def objtype(self):
0060         return self._type
0061     def payloads(self):
0062         return [H5Payload(self._payloadGroup[p],p.split('/')[-1], self._compressor) for p in self._payloadGroup]
0063     def idToPayloadNames(self):
0064         return { self._payloadGroup[p].id:p.split('/')[-1] for p in self._payloadGroup }
0065
0066 class H5Tag(object):
0067     def __init__(self, file, group, name):
0068         self._file = file
0069
0070         compressor = None
0071         compressorName = self._file.attrs['default_payload_compressor']
0072         if compressorName == 'lzma':
0073             compressor = lzma
0074         if compressorName == 'zlib':
0075             compressor = zlib
0076         self._group = group
0077         self._record = self._group.attrs['record']
0078         self._name = name
0079
0080         recordGroup = file['Records'][self._record]
0081         dataProductsGroup = recordGroup['DataProducts']
0082
0083         self._dataProducts = [H5DataProduct(dataProductsGroup[g],g.split('/')[-1], compressor) for g in dataProductsGroup]
0084         self._dbtags = self._group.attrs['db_tags']
0085         self._time_type = self._group.attrs['time_type']
0086     def record(self):
0087         return self._record
0088     def name(self):
0089         return self._name
0090     def time_type(self):
0091         return self._time_type
0092     def originalTagNames(self):
0093         return self._dbtags
0094     def iovsNPayloadNames(self):
0095
0096         #asking an h5 object for its name is a slow operation
0097         idToName = {self._file['null_payload'].id: None}
0098         for d in self._dataProducts:
0099             idToName.update(d.idToPayloadNames())
0100
0101         first = self._group['first'][()]
0102         last = self._group['last'][()]
0103         payloadRefs = self._group['payload']
0104         return list(zip( (IOVSyncValue(x['high'],x['low']) for x in first),
0105                          (IOVSyncValue(x['high'],x['low']) for x in last),
0106                          ([idToName[self._file[r].id] for r in refs] for refs in payloadRefs)) )
0107
0108     def dataProducts(self):
0109         return self._dataProducts
0110
0111 class H5GlobalTag(object):
0112     def __init__(self, filename, name):
0113         self._file = h5py.File(filename,'r')
0114         self._name = name
0115
0116     def tags(self):
0117         #looking up names is slow so better to make cache
0118         tagID2Name = {}
0119         recordsGroup = self._file['Records']
0120         for recordName in recordsGroup:
0121             r = recordsGroup[recordName]
0122             tagsGroup = r['Tags']
0123             for tagName in tagsGroup:
0124                 tagID2Name[tagsGroup[tagName].id] = tagName
0125         globalTagGroup = self._file['GlobalTags'][self._name]
0126         return (H5Tag(self._file, self._file[t], tagID2Name[self._file[t].id]) for t in globalTagGroup['Tags'])
0127
0128 def main():
0129     parser = argparse.ArgumentParser(description='Read from HDF5 file and write to HDF5 file')
0130     parser.add_argument('input', help="Name of file to read")
0131     parser.add_argument('name', nargs='+', help="Name of the global tag.")
0132
0133     parser.add_argument('--exclude', '-e', nargs='*', help = 'list of records to exclude from the file (can not be used with --include)')
0134     parser.add_argument('--include', '-i', nargs='*', help = 'lost of the only records that should be included in the file (can not be used with --exclude')
0135     parser.add_argument('--output', '-o', default='test.h5cond', help='name of hdf5 output file to write')
0136     parser.add_argument('--compressor', '-c', default='zlib', choices=['zlib', 'lzma', 'none'], help="compress data using 'zlib', 'lzma' or 'none'")
0137     args = parser.parse_args()
0138
0139     if args.exclude and args.include:
0140         print("Can not use --exclude and --include at the same time")
0141         exit(-1)
0142
0143     excludeRecords = set()
0144     if args.exclude:
0145         excludeRecords = set(args.exclude)
0146     includeRecords = set()
0147     if args.include:
0148         includeRecords = set(args.include)
0149
0150     writeH5File(args.output, args.name, excludeRecords, includeRecords, lambda x: H5GlobalTag(args.input, x),  args.compressor)
0151 if __name__ == '__main__':
0152     main()