A variable can now be created totally at once (values+attributes)

Signed-off-by: Alexis Jeandet <alexis.jeandet@member.fsf.org>

pycdfpp/__init__.py

@@ -22,7 +22,7 @@
 import numpy as np
 from ._pycdfpp import *
 from . import _pycdfpp
-from typing import ByteString
+from typing import ByteString, Mapping, List, Any
 import sys
 import os
 __here__ = os.path.dirname(os.path.abspath(__file__))
@@ -59,14 +59,15 @@
     CDF_TIME_TT2000
 )
 
+
 def _values_from_np_array(values: np.ndarray, cdf_type=None):
     if values.dtype.num == 21:
         if cdf_type in (None, CDF_TIME_TT2000, CDF_EPOCH, CDF_EPOCH16):
             return (values.view(np.uint64),
-                             cdf_type or CDF_TIME_TT2000)
+                    cdf_type or CDF_TIME_TT2000)
     else:
         return (values, cdf_type or _NUMPY_TO_CDF_TYPE_[
-                         values.dtype.num])
+            values.dtype.num])
 
 
 def _patch_set_values():
@@ -75,19 +76,29 @@ def _set_values_wrapper(self, values: np.ndarray, cdf_type=None):
 
     Variable.set_values = _set_values_wrapper
 
+
 def _patch_add_variable():
-        def _add_variable_wrapper(self,name:str, values: np.ndarray or None=None, cdf_type=None, is_nrv:bool=False,compression:CDF_compression_type=CDF_compression_type.no_compression):
-            if values is not None:
-                v,t=_values_from_np_array(values, cdf_type)
-                return self._add_variable(name=name, values=v, cdf_type=t, is_nrv=is_nrv, compression=compression)
-            else:
-                return self._add_variable(name=name, is_nrv=is_nrv, compression=compression)
+    def _add_variable_wrapper(self, name: str, values: np.ndarray or None = None, cdf_type=None, is_nrv: bool = False,
+                              compression: CDF_compression_type = CDF_compression_type.no_compression,
+                              attributes: Mapping[str, List[Any]] or None = None):
+        if values is not None:
+            v, t = _values_from_np_array(values, cdf_type)
+            var = self._add_variable(
+                name=name, values=v, cdf_type=t, is_nrv=is_nrv, compression=compression)
+        else:
+            var = self._add_variable(
+                name=name, is_nrv=is_nrv, compression=compression)
+        if attributes is not None and var is not None:
+            for name, values in attributes.items():
+                var.add_attribute(name, values)
+        return var
+    CDF.add_variable = _add_variable_wrapper
 
-        CDF.add_variable = _add_variable_wrapper
 
 _patch_set_values()
 _patch_add_variable()
 
+
 def to_datetime64(values):
     """
     to_datetime64
@@ -168,7 +179,7 @@ def to_epoch16(values):
     return _pycdfpp.to_epoch16(values)
 
 
-def load(file_or_buffer : str or ByteString, iso_8859_1_to_utf8:bool = False, lazy_load: bool =True):
+def load(file_or_buffer: str or ByteString, iso_8859_1_to_utf8: bool = False, lazy_load: bool = True):
     """
     Load and parse a CDF file.
 

pycdfpp/chrono.hpp

@@ -32,6 +32,7 @@ using namespace cdf;
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
+#include <pybind11/operators.h>
 
 namespace py = pybind11;
 
@@ -203,11 +204,14 @@ void def_time_types_wrapper(T& mod)
 {
     py::class_<tt2000_t>(mod, "tt2000_t")
         .def_readwrite("value", &tt2000_t::value)
+        .def(py::self == py::self)
         .def("__repr__", __repr__<tt2000_t>);
     py::class_<epoch>(mod, "epoch")
         .def_readwrite("value", &epoch::value)
+        .def(py::self == py::self)
         .def("__repr__", __repr__<epoch>);
     py::class_<epoch16>(mod, "epoch16")
+        .def(py::self == py::self)
         .def_readwrite("seconds", &epoch16::seconds)
         .def_readwrite("picoseconds", &epoch16::picoseconds)
         .def("__repr__", __repr__<epoch16>);

tests/python_saving/test.py

@@ -24,6 +24,18 @@ def test_can_create_CDF_attributes(self):
         cdf = pycdfpp.CDF()
         cdf.add_attribute("test_attribute", [[1,2,3], [datetime(2018,1,1), datetime(2018,1,2)], "hello\nworld"])
 
+    def test_can_create_variable_and_attributes_at_once(self):
+        cdf = pycdfpp.CDF()
+        cdf.add_variable("test_variable", attributes={"attr1":[1,2,3], "attr2":[datetime(2018,1,1), datetime(2018,1,2)] })
+        self.assertListEqual(cdf["test_variable"].attributes["attr1"][0], [1,2,3])
+        self.assertListEqual(cdf["test_variable"].attributes["attr2"][0], [pycdfpp.to_tt2000(datetime(2018,1,1)), pycdfpp.to_tt2000(datetime(2018,1,2))])
+
+    def test_can_create_an_empty_variable(self):
+        cdf = pycdfpp.CDF()
+        cdf.add_variable("test_variable", values=np.ones((0,100,10), dtype=np.float32))
+        self.assertListEqual(cdf["test_variable"].shape, [0,100,10])
+
+
     def test_can_save_an_empty_CDF_object(self):
         with NamedTemporaryFile() as f:
             self.assertTrue(pycdfpp.save(pycdfpp.CDF(),f.name))