vtu_unit_test.py

import unyt
import xmltodict
import sys
import os


class VTKUnitChecker(object):
    """a recursive unit checker for xml vtu/pvtu files from ASPECT

    This module recursively searches the elements of an xml vtk files, looking for
    data entries. Elements identified as data entries are then parsed for the
    new units tag and the resulting unit string is run through unyt to see if it
    is valid.

    This requires the xmltodict and unyt packages.

    I assumed the new unit tag will be named "Units", if that's not the case
    you can edit the code or use the `unit_tag` parameter (see example below)

    This should work on both vtu and pvtu files -- it looks for the xml
    elements containing the string 'dataarray', so it will catch both
    'PDataArray' and 'DataArray' elements. This too can be adjusted with
    the parameters (see example below).

    Example
    -------

    To use from the command line, from the folder containing this file:

        $ python vtu_unit_test.py /path/to/file.vtu /different/file.pvtu

    Running from the command line will print output like:

        testing file aspect/output_shell_simple_2d/solution/solution-00000.pvtu

            Could not find a units tag for velocity
            Could not find a units tag for p
            Could not find a units tag for T

        testing file aspect/output_shell_simple_2d/solution/solution-00001.pvtu

            velocity passed unit test with m/s
            p passed unit test with nondimensional units
            *** test failure! *** the unit string bad for T could not be handled.
                Try `import unyt; unyt.unyt_quantity(1, bad)`  from a python shell
                to see the full error message.

    To check a single file from a python shell launched from the folder containing
    this file:

        >>> from vtu_unit_test import VTKUnitChecker
        >>> fi = "/path/to/vtu_file.vtu"
        >>> VTKUnitChecker(fi).check_vtu_file()

    To change the name of the unit tag,
        >>> VTKUnitChecker(fi, unit_tag="Unit").check_vtu_file()

    Or to change the tag that identifies a data element:
        >>> VTKUnitChecker(fi, data_tag="PDataArray").check_vtu_file()


    """
    def __init__(self, vtu_file: str, unit_tag: str = 'Units', data_tag: str = 'dataarray'):

        if os.path.isfile(vtu_file) is False:
            raise OSError(f"{vtu_file} not found")

        self.vtu_file = vtu_file
        self.data_entries = []
        self.unit_tag = "@" + unit_tag
        self.data_tag = data_tag.lower()

    def check_vtu_file(self):
        with open(self.vtu_file) as vtu_fi:
            vtu_xml = xmltodict.parse(vtu_fi.read())

        # build the list of DataArray entries
        self.recursive_attr_check(vtu_xml)

        # the units of each field
        self.field_units = [DataField(data, self) for data in self.data_entries]

        # print out the results
        self.print_result()

    def print_result(self):
        if len(self.field_units) == 0:
            print("    No data fields found")
            return

        for result in self.field_units:
            print(result)

    def recursive_attr_check(self, xmldict):
        # searches a nested dictionary for DataArray keys, allowing for lists
        # DataArray entries will be appended to the data_entries list and processed
        # subsequently

        if isinstance(xmldict, dict):
            for key, val in xmldict.items():
                if self.data_tag in key.lower() and type(val) is list:
                    # will catch DataArray, PDataArray
                    self.data_entries = self.data_entries + val
                if isinstance(val, dict) or type(val) is list:
                    self.recursive_attr_check(val)
        elif type(xmldict) is list:
            for val in xmldict:
                if type(val) in [dict, list]:
                    self.recursive_attr_check(val)


class DataField:

    def __init__(self, data_entry, vtu_unit_checker):
        self.field = data_entry['@Name']
        self.has_units_tag = vtu_unit_checker.unit_tag in data_entry
        if self.has_units_tag:
            self.units = data_entry[vtu_unit_checker.unit_tag]
        else:
            self.units = None

        self.passed = True
        self.test_result = self.test_units()

    def test_units(self):

        if self.has_units_tag:
            try:
                _ = unyt.unyt_quantity(1, self.units)
            except unyt.exceptions.UnitParseError:
                self.passed = False
                return "unit NOT recognized"
            return "unit recognized"
        else:
            return "unit tag not found"

    def __str__(self):

        if self.has_units_tag:
            if self.passed:
                pstr = f"{self.field} passed unit test "
                if self.units == "":
                    pstr += "with nondimensional units"
                else:
                    pstr += f"with {self.units}"

            else:
                pstr = f"*** test failure! *** the unit string {self.units} for {self.field}" \
                       f" could not be handled. Try `import unyt; unyt.unyt_quantity(1, {self.units})` " \
                       f" from a python shell to see the full error message."
        else:
            pstr = f"Could not find a units tag for {self.field}"

        return "    " + pstr


if __name__ == "__main__":

    if len(sys.argv) < 2:
        print("please provide list of vtu files to check, e.g.:")
        print("    $ python vtu_unit_test.py /path/to/file.vtu /path/to/different/vtu")
        exit()

    for testfi in sys.argv[1:]:
        print("")
        print(f"testing file {testfi}")
        print("")
        VTKUnitChecker(testfi).check_vtu_file()