Merge pull request #4043 from pnorbert/python-fixes

Add arguments to IO.AvailableAttributes() and io.available_attributes…

bindings/Python/py11IO.cpp

@@ -344,10 +344,11 @@ std::map<std::string, Params> IO::AvailableVariables()
     return m_IO->GetAvailableVariables();
 }
 
-std::map<std::string, Params> IO::AvailableAttributes()
+std::map<std::string, Params> IO::AvailableAttributes(const std::string &varname,
+                                                      const std::string &separator)
 {
     helper::CheckForNullptr(m_IO, "in call to IO::AvailableAttributes");
-    return m_IO->GetAvailableAttributes();
+    return m_IO->GetAvailableAttributes(varname, separator);
 }
 
 std::string IO::VariableType(const std::string &name) const

bindings/Python/py11IO.h

@@ -113,7 +113,8 @@ class IO
 
     std::map<std::string, Params> AvailableVariables();
 
-    std::map<std::string, Params> AvailableAttributes();
+    std::map<std::string, Params> AvailableAttributes(const std::string &varname = "",
+                                                      const std::string &separator = "/");
 
     std::string VariableType(const std::string &name) const;
 

bindings/Python/py11glue.cpp

@@ -292,8 +292,11 @@ PYBIND11_MODULE(ADIOS2_PYTHON_MODULE_NAME, m)
                                                                 adios2::py11::MPI4PY_Comm comm)) &
                          adios2::py11::IO::Open)
 #endif
+        .def("AvailableAttributes", &adios2::py11::IO::AvailableAttributes,
+             pybind11::arg("varname") = "", pybind11::arg("separator") = "/",
+             pybind11::return_value_policy::move)
+
         .def("AvailableVariables", &adios2::py11::IO::AvailableVariables)
-        .def("AvailableAttributes", &adios2::py11::IO::AvailableAttributes)
         .def("FlushAll", &adios2::py11::IO::FlushAll)
         .def("EngineType", &adios2::py11::IO::EngineType)
         .def("RemoveVariable", &adios2::py11::IO::RemoveVariable)

python/adios2/io.py

@@ -99,19 +99,30 @@ def inquire_attribute(self, name, variable_name="", separator="/"):
             attr.impl = attrimpl
         return attr
 
-    def available_attributes(self):
+    def available_attributes(self, varname="", separator="/"):
         """
         Returns a 2-level dictionary with attribute information.
         Read mode only.
 
+        Parameters
+            variable_name
+                If varname is set, attributes assigned to that variable are returned.
+                The keys returned are attribute names with the prefix of varname + separator
+                removed.
+
+            separator
+                concatenation string between variable_name and attribute
+                e.g. varname + separator + name ("var/attr")
+                Not used if varname is empty
+
         Returns
             attributes dictionary
                 key
                     attribute name
                 value
                     attribute information dictionary
         """
-        return self.impl.AvailableAttributes()
+        return self.impl.AvailableAttributes(varname, separator)
 
     def remove_attribute(self, name):
         """

python/adios2/stream.py

@@ -172,19 +172,30 @@ def available_variables(self):
         """
         return self._io.available_variables()
 
-    def available_attributes(self):
+    def available_attributes(self, varname="", separator="/"):
         """
         Returns a 2-level dictionary with attribute information.
         Read mode only.
 
+        Parameters
+            variable_name
+                If varname is set, attributes assigned to that variable are returned.
+                The keys returned are attribute names with the prefix of varname + separator
+                removed.
+
+            separator
+                concatenation string between variable_name and attribute
+                e.g. varname + separator + name ("var/attr")
+                Not used if varname is empty
+
         Returns
             attributes dictionary
                 key
                     attribute name
                 value
                     attribute information dictionary
         """
-        return self._io.available_attributes()
+        return self._io.available_attributes(varname, separator)
 
     def define_variable(self, name):
         """

testing/adios2/bindings/python/TestBPWriteReadTypes.py

@@ -49,8 +49,8 @@ def check_array(np1, np2, hint):
     "attrU16",
     "attrU32",
     "attrU64",
-    "attrR32",
-    "attrR64",
+    "varR32/attrR32",
+    "varR64::attrR64",
 ]
 var_names = [
     "varStr",
@@ -104,8 +104,9 @@ def check_array(np1, np2, hint):
 attU16 = ioWriter.DefineAttribute("attrU16", data.U16)
 attU32 = ioWriter.DefineAttribute("attrU32", data.U32)
 attU64 = ioWriter.DefineAttribute("attrU64", data.U64)
-attR32 = ioWriter.DefineAttribute("attrR32", data.R32)
-attR64 = ioWriter.DefineAttribute("attrR64", data.R64)
+# add an attribute to a variable
+attR32 = ioWriter.DefineAttribute("attrR32", data.R32, "varR32")
+attR64 = ioWriter.DefineAttribute("attrR64", data.R64, "varR64", "::")
 
 ioWriter.SetEngine("BPFile")
 ioParams = {"Threads": "1", "InitialBufferSize": "17Kb"}
@@ -166,8 +167,8 @@ def check_array(np1, np2, hint):
 attrU16 = ioReader.InquireAttribute("attrU16")
 attrU32 = ioReader.InquireAttribute("attrU32")
 attrU64 = ioReader.InquireAttribute("attrU64")
-attrR32 = ioReader.InquireAttribute("attrR32")
-attrR64 = ioReader.InquireAttribute("attrR64")
+attrR32 = ioReader.InquireAttribute("attrR32", "varR32")
+attrR64 = ioReader.InquireAttribute("attrR64", "varR64", "::")
 
 check_object(attrString, "attrString")
 check_object(attrStringArray, "attrStringArray")
@@ -179,8 +180,14 @@ def check_array(np1, np2, hint):
 check_object(attrU16, "attrU16")
 check_object(attrU32, "attrU32")
 check_object(attrU64, "attrU64")
-check_object(attrR32, "attrR32")
-check_object(attrR64, "attrR64")
+check_object(attrR32, "varR32/attrR32")
+check_object(attrR64, "varR64::attrR64")
+
+# alternative inquire format
+attrR32 = ioReader.InquireAttribute("varR32/attrR32")
+attrR64 = ioReader.InquireAttribute("varR64::attrR64")
+check_object(attrR32, "varR32/attrR32")
+check_object(attrR64, "varR64::attrR64")
 
 attrStringData = attrString.DataString()
 print(f"attrString = {attrStringData}", flush=True)
@@ -218,6 +225,7 @@ def check_array(np1, np2, hint):
 check_array(attrR32Data, data.R32, "R32")
 check_array(attrR64Data, data.R64, "R64")
 
+print("=========== Attributes ===========")
 attributesInfo = ioReader.AvailableAttributes()
 for name, info in attributesInfo.items():
     check_name(name, attr_names)
@@ -227,6 +235,27 @@ def check_array(np1, np2, hint):
             print("\t" + key + ": " + value)
         print("\n")
 
+print("=========== Available attributes of varR32 ===========")
+attributesInfoR32 = ioReader.AvailableAttributes("varR32")
+for name, info in attributesInfoR32.items():
+    check_name(name, ["attrR32"])
+    if rank == 0:
+        print("attribute_name: " + name)
+        for key, value in info.items():
+            print("\t" + key + ": " + value)
+        print("\n")
+
+print("=========== Available attributes of varR64 ===========")
+attributesInfoR32 = ioReader.AvailableAttributes("varR64", "::")
+for name, info in attributesInfoR32.items():
+    check_name(name, ["attrR64"])
+    if rank == 0:
+        print("attribute_name: " + name)
+        for key, value in info.items():
+            print("\t" + key + ": " + value)
+        print("\n")
+
+print("=========== Variables ===========")
 varStr = ioReader.InquireVariable("varStr")
 varI8 = ioReader.InquireVariable("varI8")
 varI16 = ioReader.InquireVariable("varI16")

testing/adios2/python/TestBPWriteTypesHighLevelAPI.py

@@ -56,9 +56,8 @@
             s.write("a_float_list", data.float_list)
             s.write("a_complex_list", data.complex_list)
 
-            # single value attributes
+            # single value attributes with numpy variables
             s.write_attribute("attrStr", "Testing single string attribute")
-            s.write_attribute("attrNx", data.Nx)
             s.write_attribute("attrI8", np.array(data.i8[0]))
             s.write_attribute("attrI16", np.array(data.i16[0]))
             s.write_attribute("attrI32", np.array(data.i32[0]))
@@ -71,15 +70,19 @@
             s.write_attribute("attrR64", np.array(data.r64[0]))
             s.write_attribute("attrC64", np.array(data.c64[0]))
 
+            # single value attributes with Python variables
+            s.write_attribute("attrNx", data.Nx)
             s.write_attribute("attr_int_value", i)
             s.write_attribute("attr_float_value", f)
             s.write_attribute("attr_complex_value", c)
 
+            # array attributes with Python lists
             s.write_attribute("attr_int_list", data.int_list)
             s.write_attribute("attr_float_list", data.float_list)
             s.write_attribute("attr_complex_list", data.complex_list)
-
             s.write_attribute("attrStrArray", ["string1", "string2", "string3"])
+
+            # array attributes with numpy arrays
             s.write_attribute("attrI8Array", data.i8)
             s.write_attribute("attrI16Array", data.i16)
             s.write_attribute("attrI32Array", data.i32)
@@ -106,6 +109,9 @@
         s.write("varC64", data.c64, shape, start, count)
 
         if rank == 0 and step.current_step() == 0:
+            # attribute assigned to variable
+            s.write_attribute("size", data.Nx, "varI8")
+            s.write_attribute("size", data.Nx, "varI16", "::")
             s.write_attribute("varattrStrArray", ["varattr1", "varattr2", "varattr3"], "steps")
             s.write_attribute("varattrI8Array", data.i8, "varI8")
             s.write_attribute("varattrI16Array", data.i16, "varI16")
@@ -376,6 +382,41 @@
             if not (inR64 == data.r64).all():
                 raise ValueError("var attrR64 array read failed")
 
+            # Attributes assigned to a variable
+            sizeI8 = fr_step.read_attribute("size", "varI8")
+            sizeI16 = fr_step.read_attribute("size", "varI16", "::")
+
+            if sizeI8[0] != data.Nx:
+                raise ValueError("attribute varI8/size read failed")
+
+            if sizeI16[0] != data.Nx:
+                raise ValueError("attribute varI16::size read failed")
+
+            sizeI8 = fr_step.read_attribute("varI8/size")
+            sizeI16 = fr_step.read_attribute("varI16::size")
+
+            if sizeI8[0] != data.Nx:
+                raise ValueError("attribute varI8/size read failed")
+
+            if sizeI16[0] != data.Nx:
+                raise ValueError("attribute varI16::size read failed")
+
+            step_attrs = fr_step.available_attributes()
+            # for name, info in step_attrs.items():
+            #     print(f"attribute {name} : {info}")
+
+            step_attrs = fr_step.available_attributes("varI8")
+            # for name, info in step_attrs.items():
+            #    print(f"attribute {name} : {info}")
+            if not [*step_attrs] == ["size", "varattrI8Array"]:
+                raise ValueError("getting attributes of varI8 failed")
+
+            step_attrs = fr_step.available_attributes("varI16", "::")
+            for name, info in step_attrs.items():
+                print(f"attribute {name} : {info}")
+            if not [*step_attrs] == ["size"]:
+                raise ValueError("getting attributes of varI16 with separator :: failed")
+
         stepStr = "Step:" + str(step)
 
         instepStr = fr_step.read("steps")