Improve section array map (#437)

* Refactor section array map

Improve const handling for the section array map

* Update dependencies

* Improve module constructor

* Fix type hint

* Improve stub generator

Author: gentlegiantJGC

requirements.py

@@ -4,12 +4,12 @@
 
 AMULET_COMPILER_TARGET_REQUIREMENT = "==2.0"
 
-PYBIND11_REQUIREMENT = "==3.0.0"
-AMULET_PYBIND11_EXTENSIONS_REQUIREMENT = "~=1.2.0.0a0"
+PYBIND11_REQUIREMENT = "==3.0.1"
+AMULET_PYBIND11_EXTENSIONS_REQUIREMENT = "~=1.2.0.0a2"
 AMULET_IO_REQUIREMENT = "~=1.0"
-AMULET_UTILS_REQUIREMENT = "~=1.1.3.0a1"
+AMULET_UTILS_REQUIREMENT = "~=1.1.3.0a6"
 AMULET_ZLIB_REQUIREMENT = "~=1.0.8.0a0"
-AMULET_NBT_REQUIREMENT = "~=5.0.2.0a0"
+AMULET_NBT_REQUIREMENT = "~=5.0.2.0a2"
 NUMPY_REQUIREMENT = "~=2.0"
 
 if os.environ.get("AMULET_PYBIND11_EXTENSIONS_REQUIREMENT", None):

src/amulet/core/_amulet_core.py.cpp

@@ -34,8 +34,5 @@ void init_module(py::module m)
 
 PYBIND11_MODULE(_amulet_core, m)
 {
-    py::options options;
-    options.disable_function_signatures();
-    m.def("init", &init_module, py::doc("init(arg0: types.ModuleType) -> None"));
-    options.enable_function_signatures();
+    m.def("init", &init_module, py::arg("m"));
 }

src/amulet/core/_amulet_core.pyi

@@ -4,4 +4,4 @@ import types
 
 __all__: list[str] = ["init"]
 
-def init(arg0: types.ModuleType) -> None: ...
+def init(m: types.ModuleType) -> None: ...

src/amulet/core/chunk/component/section_array_map.cpp

@@ -153,9 +153,8 @@ void SectionArrayMap::serialise(BinaryWriter& writer) const
         get_default_array());
 
     // Write arrays
-    const auto& arrays = get_arrays();
-    writer.write_numeric<std::uint64_t>(arrays.size());
-    for (const auto& [cy, arr] : arrays) {
+    writer.write_numeric<std::uint64_t>(_arrays.size());
+    for (const auto& [cy, arr] : _arrays) {
         writer.write_numeric<std::int64_t>(cy);
         arr->serialise(writer);
     }

src/amulet/core/chunk/component/section_array_map.hpp

@@ -14,6 +14,8 @@
 #include <amulet/io/binary_reader.hpp>
 #include <amulet/io/binary_writer.hpp>
 
+#include <amulet/utils/view/map_view.hpp>
+
 #include <amulet/core/dll.hpp>
 
 namespace Amulet {
@@ -41,9 +43,10 @@ class IndexArray3D {
     AMULET_CORE_EXPORT static IndexArray3D deserialise(BinaryReader&);
 
     const SectionShape& get_shape() const { return _shape; }
-    const size_t& get_size() const { return _size; }
-    std::uint32_t* get_buffer() const { return _buffer; }
-    std::span<std::uint32_t> get_span() const { return { _buffer, _size }; }
+    size_t get_size() const { return _size; }
+    std::uint32_t* get_buffer() { return _buffer; }
+    const std::uint32_t* get_buffer() const { return _buffer; }
+    std::span<std::uint32_t> get_span() { return { _buffer, _size }; }
 };
 
 class SectionArrayMap {
@@ -84,18 +87,43 @@ class SectionArrayMap {
 
     const SectionShape& get_array_shape() const { return _array_shape; }
 
-    std::variant<std::uint32_t, std::shared_ptr<IndexArray3D>> get_default_array() const
+    std::variant<std::uint32_t, std::shared_ptr<const IndexArray3D>> get_default_array() const
+    {
+        return std::visit(
+            [](auto&& arg) -> std::variant<std::uint32_t, std::shared_ptr<const IndexArray3D>> {
+                return arg;
+            },
+            _default_array);
+    }
+
+    std::variant<std::uint32_t, std::shared_ptr<IndexArray3D>> get_default_array()
     {
         return _default_array;
     }
 
-    void set_default_array(std::variant<std::uint32_t, std::shared_ptr<IndexArray3D>> default_array)
+    void set_default_array(std::uint32_t default_array)
     {
-        validate_array_shape(default_array);
+        _default_array = default_array;
+    }
+
+    void set_default_array(std::shared_ptr<IndexArray3D> default_array)
+    {
+        validate_array_shape(*default_array);
         _default_array = std::move(default_array);
     }
 
-    const std::unordered_map<std::int64_t, std::shared_ptr<IndexArray3D>>& get_arrays() const
+    void set_default_array(const IndexArray3D& default_array)
+    {
+        validate_array_shape(default_array);
+        _default_array = std::make_shared<IndexArray3D>(default_array);
+    }
+
+    const std::unordered_map<std::int64_t, std::shared_ptr<IndexArray3D>>& get_arrays()
+    {
+        return _arrays;
+    }
+
+    const MapView<std::unordered_map<std::int64_t, std::shared_ptr<IndexArray3D>>, std::shared_ptr<const IndexArray3D>> get_arrays() const
     {
         return _arrays;
     }
@@ -107,17 +135,38 @@ class SectionArrayMap {
         return _arrays.contains(cy);
     }
 
-    std::shared_ptr<IndexArray3D> get_section(std::int64_t cy) const
+    std::shared_ptr<IndexArray3D> get_section(std::int64_t cy)
     {
         return _arrays.at(cy);
     }
 
+    std::shared_ptr<const IndexArray3D> get_section(std::int64_t cy) const
+    {
+        return _arrays.at(cy);
+    }
+
+    IndexArray3D& get_section_ref(std::int64_t cy)
+    {
+        return *_arrays.at(cy);
+    }
+
+    const IndexArray3D& get_section_ref(std::int64_t cy) const
+    {
+        return *_arrays.at(cy);
+    }
+
     void set_section(std::int64_t cy, std::shared_ptr<IndexArray3D> section)
     {
         validate_array_shape(*section);
         _arrays.insert_or_assign(cy, std::move(section));
     }
 
+    void set_section(std::int64_t cy, const IndexArray3D& section)
+    {
+        validate_array_shape(section);
+        _arrays.insert_or_assign(cy, std::make_shared<IndexArray3D>(section));
+    }
+
     AMULET_CORE_EXPORT void populate_section(std::int64_t cy);
 
     void del_section(std::int64_t cy)

src/amulet/core/chunk/component/section_array_map.py.cpp

@@ -144,7 +144,7 @@ py::module init_section_array_map(py::module m_parent)
     py::object asarray = py::module::import("numpy").attr("asarray");
     SectionArrayMap.def_property(
         "default_array",
-        [asarray](const Amulet::SectionArrayMap& self) {
+        [asarray](Amulet::SectionArrayMap& self) {
             return std::visit([asarray](auto&& arg) -> std::variant<std::uint32_t, py::array> {
                 using T = std::decay_t<decltype(arg)>;
                 if constexpr (std::is_same_v<T, std::uint32_t>) {
@@ -202,7 +202,7 @@ py::module init_section_array_map(py::module m_parent)
         &Amulet::SectionArrayMap::del_section);
     SectionArrayMap.def(
         "__getitem__",
-        [asarray](const Amulet::SectionArrayMap& self, std::int64_t cy) -> pyext::PyObjectCpp<pyext::numpy::array_t<std::uint32_t>> {
+        [asarray](Amulet::SectionArrayMap& self, std::int64_t cy) -> pyext::PyObjectCpp<pyext::numpy::array_t<std::uint32_t>> {
             try {
                 return asarray(py::cast(self.get_section(cy)));
             } catch (const std::out_of_range&) {
@@ -214,7 +214,7 @@ py::module init_section_array_map(py::module m_parent)
         &Amulet::SectionArrayMap::get_size);
     SectionArrayMap.def(
         "__iter__",
-        [](const Amulet::SectionArrayMap& self) -> pyext::collections::Iterator<std::int64_t> {
+        [](Amulet::SectionArrayMap& self) -> pyext::collections::Iterator<std::int64_t> {
             return pyext::make_iterator(pyext::detail::MapIterator(self.get_arrays()));
         },
         py::keep_alive<0, 1>());

src/amulet/core/chunk/component/section_array_map.pyi

@@ -94,12 +94,8 @@ class SectionArrayMap:
     def update(
         self,
         other: (
-            collections.abc.Mapping[
-                typing.SupportsInt, numpy.typing.NDArray[numpy.uint32]
-            ]
-            | collections.abc.Iterable[
-                tuple[typing.SupportsInt, numpy.typing.NDArray[numpy.uint32]]
-            ]
+            collections.abc.Mapping[int, numpy.typing.NDArray[numpy.uint32]]
+            | collections.abc.Iterable[tuple[int, numpy.typing.NDArray[numpy.uint32]]]
         ) = (),
         **kwargs: numpy.typing.NDArray[numpy.uint32],
     ) -> None: ...

tests/test_amulet_core/_test_amulet_core.py.cpp

@@ -31,8 +31,5 @@ void init_module(py::module m){
 }
 
 PYBIND11_MODULE(_test_amulet_core, m) {
-    py::options options;
-    options.disable_function_signatures();
-    m.def("init", &init_module, py::doc("init(arg0: types.ModuleType) -> None"));
-    options.enable_function_signatures();
+    m.def("init", &init_module, py::arg("m"));
 }

tests/test_amulet_core/_test_amulet_core.pyi

@@ -4,4 +4,4 @@ import types
 
 __all__: list[str] = ["init"]
 
-def init(arg0: types.ModuleType) -> None: ...
+def init(m: types.ModuleType) -> None: ...

tools/generate_pybind_stubs.py

@@ -7,16 +7,28 @@
 import pybind11_stubgen
 from pybind11_stubgen.structs import Identifier
 from pybind11_stubgen.parser.mixins.filter import FilterClassMembers
-from pybind11_stubgen import main as pybind11_stubgen_main
+
+
+ForwardRefPattern = re.compile(r"ForwardRef\('(?P<variable>[a-zA-Z_][a-zA-Z0-9_]*)'\)")
+
+QuotePattern = re.compile(r"'(?P<variable>[a-zA-Z_][a-zA-Z0-9_]*)'")
+
+
+def fix_value(value: str) -> str:
+    value = value.replace("NoneType", "None")
+    value = ForwardRefPattern.sub(lambda match: match.group("variable"), value)
+    value = QuotePattern.sub(lambda match: match.group("variable"), value)
+    return value
+
 
 UnionPattern = re.compile(
-    r"^(?P<variable>[a-zA-Z_][a-zA-Z0-9_]*): types\.UnionType\s*#\s*value = (?P<value>.*)$",
+    r"^(?P<variable>[a-zA-Z_][a-zA-Z0-9_]*): (types\.UnionType|typing\._UnionGenericAlias)\s*#\s*value = (?P<value>.*)$",
     flags=re.MULTILINE,
 )
 
 
 def union_sub_func(match: re.Match[str]) -> str:
-    return f'{match.group("variable")}: typing.TypeAlias = {match.group("value")}'
+    return f'{match.group("variable")}: typing.TypeAlias = {fix_value(match.group("value"))}'
 
 
 ClassVarUnionPattern = re.compile(
@@ -26,7 +38,7 @@ def union_sub_func(match: re.Match[str]) -> str:
 
 
 def class_var_union_sub_func(match: re.Match) -> str:
-    return f'{match.group("variable")}: typing.TypeAlias = {match.group("value")}'
+    return f'{match.group("variable")}: typing.TypeAlias = {fix_value(match.group("value"))}'
 
 
 VersionPattern = re.compile(r"(?P<var>[a-zA-Z0-9_].*): str = '.*?'")
@@ -87,7 +99,7 @@ def eq_sub_func(match: re.Match[str]) -> str:
 
 
 def generic_alias_sub_func(match: re.Match) -> str:
-    return f"{match.group('variable')}: typing.TypeAlias = {match.group('value')}"
+    return f'{match.group("variable")}: typing.TypeAlias = {fix_value(match.group("value"))}'
 
 
 def get_module_path(name: str) -> str: