petsc/config/query_tests.py

#!/usr/bin/env python
import fnmatch
import glob
import inspect
import os
import optparse
import pickle
import re
import sys

thisfile = os.path.abspath(inspect.getfile(inspect.currentframe()))
pdir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(thisfile)))))
sys.path.insert(0, os.path.join(pdir, 'config'))

import testparse
from gmakegentest import nameSpace


"""
  Tool for querying the tests.

  Which tests to query?  Two options:
      1. Query only the tests that are run for a given configuration.
      2. Query all of the test files in the source directory
  For #1:
     Use dataDict as written out by gmakegentest.py in $PETSC_ARCH/$TESTBASE
  For #2:
     Walk the entire tree parsing the files as we go along using testparse.
     The tree walker is simpler than what is in gmakegentest.py

  The dataDict follows that generated by testparse.  gmakegentest.py does
  further manipulations of the dataDict to handle things like for loops
  so if using #2, those modifications are not included.

  Querying:
      The dataDict dictionary is then "inverted" to create a dictionary with the
      range of field values as keys and list test names as the values.  This
      allows fast searching

"""

def query(invDict,label):
    """
    Search the keys using fnmatch to find matching names and return list with
    the results
    """
    results=[]
    if 'name' in invDict:
        return fnmatch.filter(invDict['name'],label)

    for key in invDict:
        if fnmatch.filter([key],label):
            # Do not return values with not unless label itself has not
            if label.startswith('!') and not key.startswith('!'): continue
            if not label.startswith('!') and key.startswith('!'): continue
            results += invDict[key]

    return results

def get_inverse_dictionary(dataDict,field,srcdir):
    """
    Create a dictionary with the values of field as the keys, and the name of
    the tests as the results.
    """
    invDict={}
    if field == 'name': invDict['name']=[]
    for root in dataDict:
      for exfile in dataDict[root]:
        for test in dataDict[root][exfile]:
          if test in testparse.buildkeys: continue
          defroot = testparse.getDefaultOutputFileRoot(test)
          name=nameSpace(defroot,os.path.relpath(root,srcdir))
          if field == 'name':
              invDict['name'].append(name)
              continue
          if field not in dataDict[root][exfile][test]: continue
          values=dataDict[root][exfile][test][field]

          for val in values.split():
              if val in invDict:
                  invDict[val].append(name)
              else:
                  invDict[val] = [name]
    return invDict

def get_gmakegentest_data(testdir):
    """
     Write out the dataDict into a pickle file
    """
    # This needs to be consistent with gmakegentest.py of course
    fd = open(os.path.join(testdir,'datatest.pkl'), 'rb')
    dataDict=pickle.load(fd)
    fd.close()
    return dataDict

def walktree(top):
    """
    Walk a directory tree, starting from 'top'
    """
    verbose = False
    dataDict = {}
    alldatafiles = []
    for root, dirs, files in os.walk(top, topdown=False):
        if "examples" not in root: continue
        if root == 'output': continue
        if '.dSYM' in root: continue
        if verbose: print(root)

        dataDict[root] = {}

        for exfile in files:
            # Ignore emacs files
            if exfile.startswith("#") or exfile.startswith(".#"): continue
            ext=os.path.splitext(exfile)[1]
            if ext[1:] not in ['c','cxx','cpp','cu','F90','F']: continue

            # Convenience
            fullex = os.path.join(root, exfile)
            if verbose: print('   --> '+fullex)
            dataDict[root].update(testparse.parseTestFile(fullex, 0))

    return dataDict

def do_query(use_source, startdir, srcdir, testdir, field, label):
    """
    Do the actual query
    This part of the code is placed here instead of main()
    to show how one could translate this into ipython/jupyer notebook
    commands for more advanced queries
    """
    # Get dictionary
    if use_source:
        dataDict=walktree(startdir)
    else:
        dataDict=get_gmakegentest_data(testdir)

    # Get inverse dictionary for searching
    invDict=get_inverse_dictionary(dataDict, field, srcdir)
    #print(invDict)

    # Now do query
    resList=query(invDict, label)

    # Print in flat list suitable for use by gmakefile.test
    print(' '.join(resList))

    return

def main():
    parser = optparse.OptionParser(usage="%prog [options] field match_pattern")
    parser.add_option('-s', '--startdir', dest='startdir',
                      help='Where to start the recursion if not srcdir',
                      default='')
    parser.add_option('-p', '--petsc-dir', dest='petsc_dir',
                      help='Set PETSC_DIR different from environment',
                      default=os.environ.get('PETSC_DIR'))
    parser.add_option('-a', '--petsc-arch', dest='petsc_arch',
                      help='Set PETSC_ARCH different from environment',
                      default=os.environ.get('PETSC_ARCH'))
    parser.add_option('--srcdir', dest='srcdir',
                      help='Set location of sources different from PETSC_DIR/src.  Must be full path.',
                      default='src')
    parser.add_option('-t', '--testdir', dest='testdir',
                      help='Test directory if not PETSC_ARCH/tests.  Must be full path',
                      default='tests')
    parser.add_option('-u', '--use-source', action="store_false",
                      dest='use_source',
                      help='Query all sources rather than those configured in PETSC_ARCH')

    opts, args = parser.parse_args()

    # Argument Sanity checks
    if len(args) != 2:
        parser.print_usage()
        print('Arguments: ')
        print('  field:          Field to search for; e.g., requires')
        print('                  To just match names, use "name"')
        print('  match_pattern:  Matching pattern for field; e.g., cuda')
        return

    # Process arguments and options -- mostly just paths here
    field=args[0]
    match=args[1]

    petsc_dir = opts.petsc_dir
    petsc_arch = opts.petsc_arch
    petsc_full_arch = os.path.join(petsc_dir, petsc_arch)

    if opts.srcdir == 'src':
      petsc_full_src = os.path.join(petsc_dir, 'src')
    else:
      petsc_full_src = opts.srcdir
    if opts.testdir == 'tests':
      petsc_full_test = os.path.join(petsc_full_arch, 'tests')
    else:
      petsc_full_test = opts.testdir
    if opts.startdir:
      startdir=opts.startdir=petsc_full_src
    else:
      startdir=petsc_full_src

    # Options Sanity checks
    if not os.path.isdir(petsc_dir):
        print("PETSC_DIR must be a directory")
        return

    if not opts.use_source:
        if not os.path.isdir(petsc_full_arch):
            print("PETSC_DIR/PETSC_ARCH must be a directory")
            return
        elif not os.path.isdir(petsc_full_test):
            print("Testdir must be a directory"+petsc_full_test)
            return
    else:
        if not os.path.isdir(petsc_full_src):
            print("Source directory must be a directory"+petsc_full_src)
            return

    # Do the actual query
    do_query(opts.use_source, startdir, petsc_full_src, petsc_full_test, field, match)

    return


if __name__ == "__main__":
        main()