From 12eaa62fff7dbfd6cad51634a9044c9ca8feb404 Mon Sep 17 00:00:00 2001
From: Alexis Jeandet <alexis.jeandet@member.fsf.org>
Date: Mon, 4 Dec 2023 17:20:57 +0100
Subject: [PATCH] Allow to reset attributes values

Signed-off-by: Alexis Jeandet <alexis.jeandet@member.fsf.org>
---
 include/cdfpp/attribute.hpp                   | 91 +++++++++++++++++++
 include/cdfpp/cdf-file.hpp                    |  8 ++
 include/cdfpp/cdf-io/common.hpp               | 10 +-
 .../cdfpp/cdf-io/saving/create_records.hpp    |  2 +-
 .../cdfpp/cdf-io/saving/records-saving.hpp    |  2 +-
 include/cdfpp/cdf-io/saving/saving.hpp        |  4 +-
 include/cdfpp/variable.hpp                    |  4 +-
 pycdfpp/__init__.py                           | 41 ++++++++-
 pycdfpp/attribute.hpp                         | 38 +++++++-
 pycdfpp/pycdfpp.cpp                           |  2 +
 pycdfpp/variable.hpp                          |  2 +-
 tests/python_saving/test.py                   | 15 ++-
 tests/simple_open/main.cpp                    | 41 +++++++--
 13 files changed, 234 insertions(+), 26 deletions(-)

diff --git a/include/cdfpp/attribute.hpp b/include/cdfpp/attribute.hpp
index bee5dc9..a7fc347 100644
--- a/include/cdfpp/attribute.hpp
+++ b/include/cdfpp/attribute.hpp
@@ -155,6 +155,91 @@ struct Attribute
 private:
     attr_data_t data;
 };
+
+struct VariableAttribute
+{
+    using attr_data_t = data_t;
+    std::string name;
+    VariableAttribute() = default;
+    VariableAttribute(const VariableAttribute&) = default;
+    VariableAttribute(VariableAttribute&&) = default;
+    VariableAttribute& operator=(VariableAttribute&&) = default;
+    VariableAttribute& operator=(const VariableAttribute&) = default;
+    VariableAttribute(const std::string& name, attr_data_t&& data) : name { name }
+    {
+        this->data = std::move(data);
+    }
+
+    inline bool operator==(const VariableAttribute& other) const
+    {
+        return other.name == name && other.data == data;
+    }
+
+    inline bool operator!=(const VariableAttribute& other) const { return !(*this == other); }
+
+    inline CDF_Types type() const noexcept { return data.type(); }
+
+    template <CDF_Types type>
+    [[nodiscard]] inline decltype(auto) get()
+    {
+        return data.get<type>();
+    }
+
+    template <CDF_Types type>
+    [[nodiscard]] inline decltype(auto) get() const
+    {
+        return data.get<type>();
+    }
+
+    template <typename type>
+    [[nodiscard]] inline decltype(auto) get()
+    {
+        return data.get<type>();
+    }
+
+    template <typename type>
+    [[nodiscard]] inline decltype(auto) get() const
+    {
+        return data.get<type>();
+    }
+
+    inline void swap(data_t& new_data) { std::swap(data, new_data); }
+
+    inline VariableAttribute& operator=(attr_data_t& new_data)
+    {
+        data = new_data;
+        return *this;
+    }
+
+    inline VariableAttribute& operator=(attr_data_t&& new_data)
+    {
+        data = new_data;
+        return *this;
+    }
+
+    inline data_t& operator*() { return data; }
+    inline const data_t& operator*()const { return data; }
+
+    [[nodiscard]] inline data_t& value() { return data; }
+    [[nodiscard]] inline const data_t& value() const { return data; }
+
+    template <typename... Ts>
+    friend void visit(Attribute& attr, Ts... lambdas);
+
+    template <typename... Ts>
+    friend void visit(const Attribute& attr, Ts... lambdas);
+
+    template <class stream_t>
+    inline stream_t& __repr__(stream_t& os, indent_t indent = {}) const
+    {
+        os << indent << name << ": " << data << std::endl;
+        return os;
+    }
+
+private:
+    data_t data;
+};
+
 template <typename... Ts>
 void visit(Attribute& attr, Ts... lambdas)
 {
@@ -176,3 +261,9 @@ inline stream_t& operator<<(stream_t& os, const cdf::Attribute& attribute)
 {
     return attribute.template __repr__<stream_t>(os);
 }
+
+template <class stream_t>
+inline stream_t& operator<<(stream_t& os, const cdf::VariableAttribute& attribute)
+{
+    return attribute.template __repr__<stream_t>(os);
+}
diff --git a/include/cdfpp/cdf-file.hpp b/include/cdfpp/cdf-file.hpp
index 2c3b4bf..8f2887f 100644
--- a/include/cdfpp/cdf-file.hpp
+++ b/include/cdfpp/cdf-file.hpp
@@ -35,6 +35,14 @@ inline stream_t& operator<<(stream_t& os, const cdf_map<std::string, cdf::Variab
     return os;
 }
 
+template <class stream_t>
+inline stream_t& operator<<(stream_t& os, const cdf_map<std::string, cdf::Attribute>& attributes)
+{
+    std::for_each(std::cbegin(attributes), std::cend(attributes),
+        [&os](const auto& item) { item.second.__repr__(os, indent_t {}); });
+    return os;
+}
+
 
 namespace cdf
 {
diff --git a/include/cdfpp/cdf-io/common.hpp b/include/cdfpp/cdf-io/common.hpp
index f718f2c..b980ea9 100644
--- a/include/cdfpp/cdf-io/common.hpp
+++ b/include/cdfpp/cdf-io/common.hpp
@@ -110,7 +110,7 @@ struct cdf_repr
     std::tuple<uint32_t, uint32_t, uint32_t> distribution_version;
     cdf_map<std::string, Variable> variables;
     cdf_map<std::string, Attribute> attributes;
-    std::vector<cdf_map<std::string, Attribute>> var_attributes;
+    std::vector<cdf_map<std::string, VariableAttribute>> var_attributes;
     cdf_majority majority;
     cdf_compression_type compression_type;
     bool lazy;
@@ -127,19 +127,19 @@ void add_global_attribute(cdf_repr& repr, const std::string& name, Attribute::at
 }
 
 void add_var_attribute(cdf_repr& repr, const std::vector<uint32_t>& variable_indexes,
-    const std::string& name, Attribute::attr_data_t&& data)
+    const std::string& name, std::vector<VariableAttribute::attr_data_t>&& data)
 {
     assert(std::size(data) == std::size(variable_indexes));
-    cdf_map<uint32_t, cdf_map<std::string, Attribute::attr_data_t>> storage;
+    cdf_map<uint32_t, cdf_map<std::string, VariableAttribute::attr_data_t>> storage;
     for (auto index = 0UL; index < std::size(data); index++)
     {
-        storage[variable_indexes[index]][name].push_back(data[index]);
+        storage[variable_indexes[index]][name] = data[index];
     }
     for (auto& [v_index, attr] : storage)
     {
         for (auto& [attr_name, attr_data] : attr)
         {
-            repr.var_attributes[v_index][attr_name] = Attribute { attr_name, std::move(attr_data) };
+            repr.var_attributes[v_index][attr_name] = VariableAttribute { attr_name, std::move(attr_data) };
         }
     }
 }
diff --git a/include/cdfpp/cdf-io/saving/create_records.hpp b/include/cdfpp/cdf-io/saving/create_records.hpp
index d106673..c45d7c2 100644
--- a/include/cdfpp/cdf-io/saving/create_records.hpp
+++ b/include/cdfpp/cdf-io/saving/create_records.hpp
@@ -127,7 +127,7 @@ namespace saving
             auto& vac = svg_ctx.body.variable_attributes[name];
             update_size(vac.adr);
             vac.attrs.push_back(&attribute);
-            const auto& data = attribute[0UL];
+            const auto& data = *attribute;
             auto& aedr = vac.aedrs.emplace_back(cdf_AzEDR_t<v3x_tag> { {}, 0, vac.adr.record.num,
                 data.type(), static_cast<int32_t>(variable.number), 0, 0, 0, 0, 0, 0 });
             if (is_string(data.type()))
diff --git a/include/cdfpp/cdf-io/saving/records-saving.hpp b/include/cdfpp/cdf-io/saving/records-saving.hpp
index 1eb2755..a0e8409 100644
--- a/include/cdfpp/cdf-io/saving/records-saving.hpp
+++ b/include/cdfpp/cdf-io/saving/records-saving.hpp
@@ -91,7 +91,7 @@ struct file_attribute_ctx
 struct variable_attribute_ctx
 {
     int32_t number;
-    std::vector<const Attribute*> attrs;
+    std::vector<const VariableAttribute*> attrs;
     record_wrapper<cdf_ADR_t<v3x_tag>> adr;
     std::vector<record_wrapper<cdf_AzEDR_t<v3x_tag>>> aedrs;
 };
diff --git a/include/cdfpp/cdf-io/saving/saving.hpp b/include/cdfpp/cdf-io/saving/saving.hpp
index ec09275..895748f 100644
--- a/include/cdfpp/cdf-io/saving/saving.hpp
+++ b/include/cdfpp/cdf-io/saving/saving.hpp
@@ -86,7 +86,7 @@ namespace saving
     }
 
     template <typename U>
-    void write_records(const std::vector<const Attribute*> attrs,
+    void write_records(const std::vector<const VariableAttribute*> attrs,
         const std::vector<record_wrapper<cdf_AzEDR_t<v3x_tag>>>& aedrs, U&& writer,
         std::size_t virtual_offset = 0)
     {
@@ -96,7 +96,7 @@ namespace saving
             auto& aedr = aedrs[i];
             auto& attr = *attrs[i];
             save_record(aedr.record, writer);
-            const auto& values = attr[0];
+            const auto& values = *attr;
             auto offset = writer.write(values.bytes_ptr(), values.bytes()) + virtual_offset;
             assert(offset - aedr.size == aedr.offset);
         }
diff --git a/include/cdfpp/variable.hpp b/include/cdfpp/variable.hpp
index de46fd8..dce0db3 100644
--- a/include/cdfpp/variable.hpp
+++ b/include/cdfpp/variable.hpp
@@ -34,7 +34,7 @@
 #include <vector>
 
 template <class stream_t>
-inline stream_t& operator<<(stream_t& os, const cdf_map<std::string, cdf::Attribute>& attributes)
+inline stream_t& operator<<(stream_t& os, const cdf_map<std::string, cdf::VariableAttribute>& attributes)
 {
     std::for_each(std::cbegin(attributes), std::cend(attributes),
         [&os](const auto& item) { item.second.__repr__(os, indent_t {}); });
@@ -68,7 +68,7 @@ struct Variable
 {
     using var_data_t = data_t;
     using shape_t = no_init_vector<uint32_t>;
-    cdf_map<std::string, Attribute> attributes;
+    cdf_map<std::string, VariableAttribute> attributes;
     Variable() = default;
     Variable(Variable&&) = default;
     Variable(const Variable&) = default;
diff --git a/pycdfpp/__init__.py b/pycdfpp/__init__.py
index a873044..82559e2 100644
--- a/pycdfpp/__init__.py
+++ b/pycdfpp/__init__.py
@@ -16,7 +16,7 @@
 
 import numpy as np
 #from ._pycdfpp import *
-from ._pycdfpp import DataType, CompressionType, Majority, Variable, Attribute, CDF, tt2000_t, epoch, epoch16, save
+from ._pycdfpp import DataType, CompressionType, Majority, Variable, VariableAttribute, Attribute, CDF, tt2000_t, epoch, epoch16, save
 from datetime import datetime
 from . import _pycdfpp
 from typing import ByteString, Mapping, List, Any
@@ -324,11 +324,50 @@ def _add_attribute_wrapper(self, name: str, entries_values: List[np.ndarray or L
         return self._add_attribute(name=name, entries_values=v, entries_types=t)
     CDF.add_attribute = _add_attribute_wrapper
 
+def _patch_attribute_set_values():
+    def _attribute_set_values(self, entries_values: List[np.ndarray or List[float or int or datetime] or str], entries_types: List[DataType or None] or None = None):
+        """Sets the values of the attribute with the given values.
+        If entries_types is not None, the attribute is created with the given data type. Otherwise, the data type is inferred from the values.
+
+        Parameters
+        ----------
+        entries_values : List[np.ndarray or List[float or int or datetime] or str]
+            The values entries to set for the attribute.
+            When a list is passed, the values are converted to a numpy.ndarray with the appropriate data type, with integers, it will choose the smallest data type that can hold all the values.
+        entries_types : List[DataType] or None, optional
+            The data type for each entry of the attribute. If None, the data type is inferred from the values. (Default is None)
+
+        """
+        entries_types = entries_types or [None]*len(entries_values)
+        v, t = [list(l) for l in zip(*[_attribute_values_view_and_type(values, data_type)
+                                       for values, data_type in zip(entries_values, entries_types)])]
+        self._set_values(v, t)
+    Attribute.set_values = _attribute_set_values
+
+def _patch_var_attribute_set_value():
+    def _attribute_set_value(self, value: np.ndarray or List[float or int or datetime] or str, data_type=None):
+        """Sets the value of the attribute with the given value.
+        If data_type is not None, the attribute is created with the given data type. Otherwise, the data type is inferred from the values.
+
+        Parameters
+        ----------
+        values : np.ndarray or List[float or int or datetime] or str
+            The values to set for the attribute.
+            When a list is passed, the values are converted to a numpy.ndarray with the appropriate data type, with integers, it will choose the smallest data type that can hold all the values.
+        data_type : DataType or None, optional
+            The data type of the attribute. If None, the data type is inferred from the values. (Default is None)
+
+        """
+        v, t = _attribute_values_view_and_type(value, data_type)
+        self._set_value(v, t)
+    VariableAttribute.set_value = _attribute_set_value
 
 _patch_add_cdf_attribute()
 _patch_add_variable_attribute()
 _patch_set_values()
 _patch_add_variable()
+_patch_attribute_set_values()
+_patch_var_attribute_set_value()
 
 
 def to_datetime64(values):
diff --git a/pycdfpp/attribute.hpp b/pycdfpp/attribute.hpp
index d86147e..1ee944c 100644
--- a/pycdfpp/attribute.hpp
+++ b/pycdfpp/attribute.hpp
@@ -356,14 +356,14 @@ Attribute::attr_data_t to_attr_data_entries(
     }
 }
 
-[[nodiscard]] Attribute& add_attribute(
+[[nodiscard]] VariableAttribute& add_attribute(
     Variable& var, const std::string& name, const string_or_buffer_t& values, CDF_Types cdf_types)
 {
     auto [it, success] = var.attributes.emplace(
-        name, name, Attribute::attr_data_t { to_attr_data_entry(values, cdf_types) });
+        name, name, VariableAttribute::attr_data_t { to_attr_data_entry(values, cdf_types) });
     if (success)
     {
-        Attribute& attr = it->second;
+        VariableAttribute& attr = it->second;
         return attr;
     }
     else
@@ -372,6 +372,17 @@ Attribute::attr_data_t to_attr_data_entries(
     }
 }
 
+void set_vattr_value(
+    VariableAttribute& attr, const string_or_buffer_t& value, CDF_Types cdf_type)
+{
+    attr = to_attr_data_entry(value, cdf_type);
+}
+
+void set_attr_values(Attribute& attr, const std::vector<string_or_buffer_t>& values,
+    const std::vector<CDF_Types>& cdf_types)
+{
+    attr = to_attr_data_entries(values, cdf_types);
+}
 
 template <typename T>
 void def_attribute_wrapper(T& mod)
@@ -381,6 +392,7 @@ void def_attribute_wrapper(T& mod)
         .def(py::self == py::self)
         .def(py::self != py::self)
         .def("__repr__", __repr__<Attribute>)
+        .def("_set_values", set_attr_values, py::arg("values").noconvert(), py::arg("data_type"))
         .def(
             "__getitem__",
             [](Attribute& att, std::size_t index) -> py_cdf_attr_data_t
@@ -400,4 +412,24 @@ void def_attribute_wrapper(T& mod)
                         "Trying to get an attribute value outside of its range");
                 return att[index].type();
             });
+
+    py::class_<VariableAttribute>(mod, "VariableAttribute")
+        .def_property_readonly("name", [](VariableAttribute& attr) { return attr.name; })
+        .def(py::self == py::self)
+        .def(py::self != py::self)
+        .def("__repr__", __repr__<VariableAttribute>)
+        .def("_set_value", set_vattr_value, py::arg("value").noconvert(), py::arg("data_type"))
+        .def(
+            "__getitem__",
+            [](VariableAttribute& att, std::size_t index) -> py_cdf_attr_data_t
+            {
+                if (index !=0 )
+                    throw std::out_of_range(
+                        "Trying to get an attribute value outside of its range");
+                return to_py_cdf_data(*att);
+            },
+            py::return_value_policy::copy)
+        .def("__len__", [](const VariableAttribute& ) { return 1; })
+        .def_property_readonly("value", [](VariableAttribute& att) -> py_cdf_attr_data_t { return to_py_cdf_data(*att); })
+        .def("type", [](VariableAttribute& att) { return att.type(); });
 }
diff --git a/pycdfpp/pycdfpp.cpp b/pycdfpp/pycdfpp.cpp
index c45c367..9ff1628 100644
--- a/pycdfpp/pycdfpp.cpp
+++ b/pycdfpp/pycdfpp.cpp
@@ -45,6 +45,7 @@ using namespace cdf;
 #include <fmt/ranges.h>
 
 PYBIND11_MAKE_OPAQUE(cdf_map<std::string, Attribute>);
+PYBIND11_MAKE_OPAQUE(cdf_map<std::string, VariableAttribute>);
 PYBIND11_MAKE_OPAQUE(cdf_map<std::string, Variable>);
 
 namespace py = pybind11;
@@ -102,6 +103,7 @@ PYBIND11_MODULE(_pycdfpp, m)
 
     def_cdf_map<std::string, Variable>(m, "VariablesMap");
     def_cdf_map<std::string, Attribute>(m, "AttributeMap");
+    def_cdf_map<std::string, VariableAttribute>(m, "VariableAttributeMap");
 
 
     def_attribute_wrapper(m);
diff --git a/pycdfpp/variable.hpp b/pycdfpp/variable.hpp
index 1155cce..5c9dfb4 100644
--- a/pycdfpp/variable.hpp
+++ b/pycdfpp/variable.hpp
@@ -242,7 +242,7 @@ void def_variable_wrapper(T& mod)
         .def_property_readonly("values_encoded", make_values_view<true>, py::keep_alive<0, 1>())
         .def("_set_values", set_values, py::arg("values").noconvert(), py::arg("data_type"))
         .def("_add_attribute",
-            static_cast<Attribute& (*)(Variable&, const std::string&, const string_or_buffer_t&,
+            static_cast<VariableAttribute& (*)(Variable&, const std::string&, const string_or_buffer_t&,
                 CDF_Types)>(add_attribute),
             py::arg { "name" }, py::arg { "values" }, py::arg { "data_type" },
             py::return_value_policy::reference_internal);
diff --git a/tests/python_saving/test.py b/tests/python_saving/test.py
index 7751a31..26f9d33 100755
--- a/tests/python_saving/test.py
+++ b/tests/python_saving/test.py
@@ -33,10 +33,23 @@ def test_compare_identical_cdfs(self):
     def test_compare_differents_cdfs(self):
         self.assertEqual(make_cdf(),pycdfpp.CDF())
 
-    def test_inmemory_save_load_empty_CDF_object(self):
+    def test_in_memory_save_load_empty_CDF_object(self):
         cdf = pycdfpp.CDF()
         self.assertIsNotNone(pycdfpp.load(pycdfpp.save(cdf)))
 
+    def test_overwrite_attribute(self):
+        cdf = make_cdf()
+        self.assertEqual(cdf.attributes["test_attribute"][0], [1,2,3])
+        self.assertEqual(cdf.attributes["test_attribute"][2], "hello\nworld")
+        cdf.attributes["test_attribute"].set_values(["hello\nworld", [datetime(2018,1,1), datetime(2018,1,2)], [1,2,3]])
+        self.assertEqual(cdf.attributes["test_attribute"][0], "hello\nworld")
+        self.assertEqual(cdf.attributes["test_attribute"][2], [1,2,3])
+
+        cdf["test_variable"].attributes["attr1"].set_value([3,2,1])
+        self.assertEqual(cdf["test_variable"].attributes["attr1"][0], [3,2,1])
+        self.assertEqual(cdf["test_variable"].attributes["attr1"].value, [3,2,1])
+
+
     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"])
diff --git a/tests/simple_open/main.cpp b/tests/simple_open/main.cpp
index 72c69b4..6420d27 100644
--- a/tests/simple_open/main.cpp
+++ b/tests/simple_open/main.cpp
@@ -56,6 +56,14 @@ std::enable_if_t<std::is_same_v<value_type, const char*>, bool> impl_compare_att
     return std::string { attr_value.data(), std::size(attr_value) } == std::string { value };
 }
 
+template <int index, typename value_type>
+std::enable_if_t<std::is_same_v<value_type, const char*>, bool> impl_compare_attribute_value(
+    const cdf::VariableAttribute& attribute, const value_type& value)
+{
+    auto attr_value = attribute.get<char>();
+    return std::string { attr_value.data(), std::size(attr_value) } == std::string { value };
+}
+
 template <int index, typename value_type>
 std::enable_if_t<std::is_scalar_v<value_type> && !std::is_same_v<value_type, const char*>, bool>
 impl_compare_attribute_value(const cdf::Attribute& attribute, const value_type& value)
@@ -64,30 +72,45 @@ impl_compare_attribute_value(const cdf::Attribute& attribute, const value_type&
 }
 
 template <int index, typename value_type>
-auto impl_compare_attribute_value(const cdf::Attribute& attribute, const value_type& value)
-    -> decltype(std::cbegin(value), value.at(0), true)
+std::enable_if_t<std::is_scalar_v<value_type> && !std::is_same_v<value_type, const char*>, bool>
+impl_compare_attribute_value(const cdf::VariableAttribute& attribute, const value_type& value)
+{
+    return attribute.get<value_type>() == no_init_vector<value_type> { value };
+}
+
+template <int index,typename attr_t, typename value_type>
+auto impl_compare_attribute_value(const attr_t& attribute, const value_type& value)
+    -> decltype(std::cbegin(value), value.at(0),attribute.value(), true)
+{
+    return attribute.template get<typename std::remove_const_t<std::remove_reference_t<value_type>>::value_type>()
+        == value;
+}
+
+template <int index,typename attr_t, typename value_type>
+auto impl_compare_attribute_value(const attr_t& attribute, const value_type& value)
+    -> decltype(std::cbegin(value), value.at(0), attribute[0], true)
 {
     return attribute
-               .get<typename std::remove_const_t<std::remove_reference_t<value_type>>::value_type>(
+               .template get<typename std::remove_const_t<std::remove_reference_t<value_type>>::value_type>(
                    index)
         == value;
 }
 
-template <int index, typename T>
-bool compare_attribute_value(const cdf::Attribute& attribute, const T& values)
+template <int index,typename attr_t,  typename T>
+bool compare_attribute_value(const attr_t& attribute, const T& values)
 {
     auto value = std::get<index>(values);
     return impl_compare_attribute_value<index>(attribute, value);
 }
 
-template <typename T, std::size_t... I>
-bool impl_compare_attribute_values(cdf::Attribute& attribute, T values, std::index_sequence<I...>)
+template <typename attr_t, typename T, std::size_t... I>
+bool impl_compare_attribute_values(attr_t& attribute, T values, std::index_sequence<I...>)
 {
     return (... && compare_attribute_value<I>(attribute, values));
 }
 
-template <typename... Ts>
-bool compare_attribute_values(cdf::Attribute& attribute, Ts... values)
+template <typename attr_t, typename... Ts>
+bool compare_attribute_values(attr_t& attribute, Ts... values)
 {
     return impl_compare_attribute_values(
         attribute, std::make_tuple(values...), std::make_index_sequence<sizeof...(values)> {});