__init__.py

 
import argparse
import basf2.utils as b2utils
 
 
class PythonVariableManager:
    '''
    Wrapper around the variable manager class.
    This is necessary for avoiding to import ROOT globally when 'variables' is imported.
    '''
 
    @classmethod
    def _instance(self):
        '''
        Return an instance to the variable manager.
        '''
        # Always avoid the top-level 'import ROOT'.
        import ROOT  # noqa
        instance = ROOT.Belle2.Variable.Manager.Instance()
        return instance
 
    def getVariable(self, *args):
        '''
        Wrapper around Manager::getVariable(std::string name) and
        Manager::getVariable(const std::string& functionName, const std::vector<std::string>& functionArguments).
        '''
        instance = PythonVariableManager._instance()
        return instance.getVariable(*args)
 
    def getVariables(self, *args):
        '''
        Wrapper around Manager::getVariables(const std::vector<std::string>& variables) and
        Manager::getVariables().
        '''
        instance = PythonVariableManager._instance()
        return instance.getVariables(*args)
 
    def addAlias(self, alias, variable):
        '''
        Wrapper around Manager::addAlias(const std::string& alias, const std::string& variable).
        '''
        instance = PythonVariableManager._instance()
        if not instance.addAlias(alias, variable):
            raise RuntimeError(f"Failed to add alias '{alias}' for variable '{variable}'")
        return alias
 
    def printAliases(self):
        '''
        Wrapper around Manager::printAliases().
        '''
        instance = PythonVariableManager._instance()
        instance.printAliases()
 
    def clearAliases(self):
        '''
        Wrapper around Manager::clearAliases().
        '''
        instance = PythonVariableManager._instance()
        instance.clearAliases()
 
    def resolveAlias(self, alias):
        '''
        Wrapper around Manager::resolveAlias(const std::string& alias).
        '''
        instance = PythonVariableManager._instance()
        return instance.resolveAlias(alias)
 
    def addCollection(self, *args):
        '''
        Wrapper around Manager::addCollection(const std::string& collection, const std::vector<std::string>& variables)
        '''
        instance = PythonVariableManager._instance()
        if not instance.addCollection(*args):
            raise RuntimeError(f"Failed to add collection with args: {args}")
 
    def getCollection(self, collection):
        '''
        Wrapper around Manager::getCollection(const std::string& collection).
        '''
        instance = PythonVariableManager._instance()
        return instance.getCollection(collection)
 
    def resolveCollections(self, *args):
        '''
        Wrapper around Manager::resolveCollections(const std::vector<std::string>& variables).
        '''
        instance = PythonVariableManager._instance()
        return instance.resolveCollections(*args)
 
    def checkDeprecatedVariable(self, variable):
        '''
        Wrapper around Manager::checkDeprecatedVariable(const std::string& name).
        '''
        instance = PythonVariableManager._instance()
        instance.checkDeprecatedVariable(variable)
 
    def evaluate(self, variable, particle):
        '''
        Wrapper around Manager::evaluate(const std::string& varName, const Particle* p).
        '''
        instance = PythonVariableManager._instance()
        return instance.evaluate(variable, particle)
 
    def evaluateVariables(self, variables, plist):
        '''
        Wrapper around Manager::evaluateVariables(const std::vector<std::string>& varNames, const ParticleList* plist).
        '''
        instance = PythonVariableManager._instance()
        return instance.evaluateVariables(variables, plist.obj())
 
    def getNames(self):
        '''
        Wrapper around Manager::getNames().
        '''
        instance = PythonVariableManager._instance()
        return instance.getNames()
 
    def getAliasNames(self):
        '''
        Wrapper around Manager::getAliasNames().
        '''
        instance = PythonVariableManager._instance()
        return instance.getAliasNames()
 
    def assertValidName(self, variable):
        '''
        Wrapper around Manager::assertValidName(const std::string& name).
        '''
        instance = PythonVariableManager._instance()
        instance.assertValidName(variable)
 
 

variables = PythonVariableManager()
 
 
def std_vector(*args):
    """
    Creates an std::vector which can be passed to pyROOT
    """
    # Always avoid the top-level 'import ROOT'.
    import ROOT  # noqa
    v = ROOT.std.vector(type(args[0]))()
    for x in args:
        v.push_back(x)
    return v
 
 
def getCommandLineOptions():
    """ Parses the command line options and returns the corresponding arguments. """
    parser = argparse.ArgumentParser()
    parser.add_argument('--no-pager', dest='pager', default=True, action='store_false',
                        help='Use a pager to show output or print to terminal.')
    args = parser.parse_args()
    return args
 
 
def printVars(changedVariableNames=None, listVars=None):
    """
    Print list of all available variables.
    """
 
    if changedVariableNames:
        print(changedVariableNames)
    print('Available variables in Variable::Manager:')
 
    if listVars is None:
        listVars = variables.getVariables()
    else:
        listVars = [variables.getVariable(name) for name in listVars]
 
    dict_VariableDataType = {0: 'double', 1: 'int', 2: 'bool'}
 
    vars = []
    for v in listVars:
        vars.append((v.group, v.name, v.description,
                     dict_VariableDataType[v.variabletype]))
 
    rows = []
    current_group = ''
    for (group, name, description, vartype) in sorted(vars):
        if current_group != group:
            current_group = group
            rows.append([group])
        rows.append([name, description, vartype])
    b2utils.pretty_print_description_list(rows)
 
 
def printVariableType(listVars=None):
 
    if isinstance(listVars, str):
        listVars = [listVars]
 
    dict_VariableDataType = {0: 'double', 1: 'int', 2: 'bool'}
 
    for var_name in listVars:
        var = variables.getVariable(var_name)
        print(f"{var_name}: {dict_VariableDataType[var.variabletype]}")
 
 
def getAllTrgNames():
    """
    Return all PSNM trigger bit names
    """
    # Always avoid the top-level 'import ROOT'.
    import ROOT  # noqa
    bits = ROOT.Belle2.PyDBObj('TRGGDLDBFTDLBits')
    evt = ROOT.Belle2.EventMetaData()
    ROOT.Belle2.DBStore.Instance().update(evt)
    size = ROOT.Belle2.TRGSummary.c_ntrgWords * ROOT.Belle2.TRGSummary.c_trgWordSize
    return [bits.getoutbitname(i) for i in range(size) if bits.getoutbitname(i) != '']