io/podio/make_datamodel_glue.py

0001 #!/usr/bin/env python3
0002 #
0003 # Copyright 2022, David Lawrence
0004 # Subject to the terms in the LICENSE file found in the top-level directory.
0005 #
0006 # This is a stop gap and not intended for long term.  2022-07-09  DL
0007 #
0008 # This will scan the list of files in the EDM4hep datamodel directory
0009 # (pointed to by the EDM4HEP_ROOT environment variable). Using the
0010 # filenames, it will generate some C++ code that can be used by
0011 # the JEventSourcePODIO and EDM4hepWriter classes to read and write all
0012 # of those types.
0013
0014 import os
0015 import sys
0016 import glob
0017
0018 print('Generating datamodel_glue.h ...')
0019
0020 # Default to "not found"
0021 WORKING_DIR = None
0022 EDM4HEP_INCLUDE_DIR = None
0023 EDM4EIC_INCLUDE_DIR = None
0024
0025 # Try getting from environment first so we can overwrite
0026 # with command line below if available.
0027 EDM4HEP_ROOT = os.environ.get("EDM4HEP_ROOT")
0028 if EDM4HEP_ROOT :
0029     EDM4HEP_INCLUDE_DIR=EDM4HEP_ROOT+'/include'
0030
0031 EDM4EIC_ROOT = os.environ.get("EDM4EIC_ROOT")
0032 if EDM4EIC_ROOT :
0033     EDM4EIC_INCLUDE_DIR=EDM4EIC_ROOT+'/include'
0034
0035
0036 # poor man's command line parsing
0037 for arg in sys.argv:
0038     if arg.startswith('WORKING_DIR'):
0039         if '=' in arg: WORKING_DIR = arg.split('=',1)[1]
0040     if arg.startswith('EDM4HEP_INCLUDE_DIR'):
0041         if '=' in arg: EDM4HEP_INCLUDE_DIR = arg.split('=',1)[1]
0042     if arg.startswith('EDM4EIC_INCLUDE_DIR'):
0043         if '=' in arg: EDM4EIC_INCLUDE_DIR = arg.split('=',1)[1]
0044
0045 # Check if EDM4HEP_ROOT is set
0046 if not EDM4HEP_INCLUDE_DIR:
0047     print("ERROR: EDM4HEP_INCLUDE_DIR not specified on command line (with \n"
0048           "EDM4HEP_INCLUDE_DIR=/path/to/edm4hep/include) and EDM4HEP_ROOT\n"
0049           "env. variable is None or empty\n"
0050           "       Please specify the EDM4HEP_INCLUDE_DIR value explicitly\n"
0051           "       or point EDM4HEP_ROOT envar to edm4hep installation root.\n"
0052           "       This script looks for '{EDM4HEP_INCLUDE_DIR}/edm4hep/*Collection.h'\n"
0053           "                          or '{EDM4HEP_ROOT}/include/edm4hep/*Collection.h'\n")
0054     sys.exit(1)
0055
0056 # Check if EDM4EIC_ROOT is set
0057 if not EDM4EIC_INCLUDE_DIR:
0058     print("ERROR: EDM4EIC_INCLUDE_DIR not specified on command line (with \n"
0059           "EDM4EIC_INCLUDE_DIR=/path/to/EDM4EIC/include) and EDM4EIC_ROOT\n"
0060           "env. variable is None or empty\n"
0061           "       Please specify the EDM4EIC_INCLUDE_DIR value explicitly\n"
0062           "       or point EDM4EIC_ROOT envar to EDM4EIC installation root.\n"
0063           "       This script looks for '{EDM4EIC_INCLUDE_DIR}/EDM4EIC/*Collection.h'\n"
0064           "                          or '{EDM4EIC_ROOT}/include/EDM4EIC/*Collection.h'\n")
0065     sys.exit(1)
0066
0067
0068 def AddCollections(datamodelName, collectionfiles):
0069     for f in collectionfiles:
0070         header_fname = f.split('/'+datamodelName)[-1]
0071         basename = header_fname.split('/')[-1].split('Collection.h')[0]
0072
0073         header = '#include <'+ datamodelName + header_fname + '>'
0074         header_lines.append(header)
0075
0076         type_map.append('namespace ' + datamodelName + ' {')
0077         type_map.append('    class ' + basename + ';')
0078         type_map.append('    class ' + basename + 'Collection;')
0079         type_map.append('    class Mutable' + basename + ';')
0080         type_map.append('};')
0081         type_map.append('#if podio_VERSION < PODIO_VERSION(0, 17, 0)')
0082         type_map.append('template <> struct PodioTypeMap<' + datamodelName + '::' + basename + '> {')
0083         type_map.append('    using collection_t = ' + datamodelName + '::' + basename + 'Collection;')
0084         type_map.append('    using mutable_t = ' + datamodelName + '::Mutable' + basename + ';')
0085         type_map.append('};')
0086         type_map.append('#endif')
0087
0088         visitor.append('        if (podio_typename == "' + datamodelName + '::' + basename + 'Collection") {')
0089         visitor.append('            return visitor(*reinterpret_cast<const ' + datamodelName + '::' + basename + 'Collection*>(&collection));')
0090         visitor.append('        }')
0091
0092
0093 collectionfiles_edm4hep = glob.glob(EDM4HEP_INCLUDE_DIR+'/edm4hep/*Collection.h')
0094 collectionfiles_edm4eic    = glob.glob(EDM4EIC_INCLUDE_DIR+'/edm4eic/*Collection.h')
0095 header_lines      = []
0096 type_map = []
0097 visitor = []
0098 AddCollections('edm4hep', collectionfiles_edm4hep)
0099 AddCollections('edm4eic'   , collectionfiles_edm4eic   )
0100
0101
0102 if WORKING_DIR : os.chdir( WORKING_DIR )
0103
0104 with open('datamodel_includes.h', 'w') as f:
0105     f.write('\n// This file automatically generated by the make_datamodel.py script\n\n')
0106     f.write('\n// IWYU pragma: begin_exports\n')
0107     f.write('\n'.join(header_lines))
0108     f.write('\n// IWYU pragma: end_exports\n')
0109     f.write('\n')
0110     f.close()
0111
0112 with open('datamodel_glue.h', 'w') as f:
0113     f.write('\n// This file automatically generated by the make_datamodel.py script\n')
0114     f.write('#pragma once\n')
0115     f.write('\n')
0116     f.write('#include <stdexcept>\n')
0117     f.write('#include <podio/podioVersion.h>\n')
0118     f.write('#include <podio/CollectionBase.h>\n')
0119     f.write('\n')
0120
0121     f.write('\ntemplate <typename T> struct PodioTypeMap {')
0122     f.write('\n#if podio_VERSION >= PODIO_VERSION(0, 17, 0)')
0123     f.write('\n    using collection_t = typename T::collection_type;')
0124     f.write('\n    using mutable_t = typename T::mutable_type;')
0125     f.write('\n#endif')
0126     f.write('\n};')
0127     f.write('\n\n')
0128     f.write('\n'.join(type_map))
0129     f.write('\n')
0130     f.write('\ntemplate <typename Visitor> struct VisitPodioCollection {')
0131     f.write('\n    void operator()(Visitor& visitor, const podio::CollectionBase& collection) {')
0132     f.write('\n        auto podio_typename = collection.getTypeName();\n')
0133     f.write('\n'.join(visitor))
0134     f.write('\n        throw std::runtime_error("Unrecognized podio typename!");')
0135     f.write('\n    }')
0136     f.write('\n};\n')
0137     f.close()