replace remaining uses of implicit thread pool

Author: lonvia

lib/simple_writer.cc

@@ -2,7 +2,7 @@
  *
  * This file is part of pyosmium. (https://osmcode.org/pyosmium/)
  *
- * Copyright (C) 2024 Sarah Hoffmann <[email protected]> and others.
+ * Copyright (C) 2025 Sarah Hoffmann <[email protected]> and others.
  * For a full list of authors see the git log.
  */
 #include <pybind11/pybind11.h>
@@ -14,10 +14,12 @@
 #include <osmium/io/header.hpp>
 #include <osmium/memory/buffer.hpp>
 #include <osmium/builder/osm_object_builder.hpp>
+#include <osmium/thread/pool.hpp>
 
 #include "cast.h"
 #include "osm_base_objects.h"
 #include "base_handler.h"
+#include "io.h"
 
 #include <filesystem>
 
@@ -30,20 +32,10 @@ class SimpleWriter : public pyosmium::BaseHandler
     enum { BUFFER_WRAP = 4096 };
 
 public:
-    SimpleWriter(const char* filename, size_t bufsz, osmium::io::Header const *header,
-                 bool overwrite, const std::string &filetype)
-    : writer(osmium::io::File(filename, filetype),
-             header ? *header : osmium::io::Header(),
-             overwrite ? osmium::io::overwrite::allow : osmium::io::overwrite::no),
-      buffer(bufsz < 2 * BUFFER_WRAP ? 2 * BUFFER_WRAP : bufsz,
-             osmium::memory::Buffer::auto_grow::yes),
-      buffer_size(buffer.capacity()) // same rounding to BUFFER_WRAP
-    {}
-
-    SimpleWriter(osmium::io::File file, size_t bufsz, osmium::io::Header const *header,
-                 bool overwrite)
-    : writer(file, header ? *header : osmium::io::Header(),
-             overwrite ? osmium::io::overwrite::allow : osmium::io::overwrite::no),
+    SimpleWriter(osmium::io::File file, unsigned long bufsz,
+                 osmium::io::Header const *header, bool overwrite,
+                 osmium::thread::Pool *pool)
+    : writer(file, header, overwrite, pool),
       buffer(bufsz < 2 * BUFFER_WRAP ? 2 * BUFFER_WRAP : bufsz,
              osmium::memory::Buffer::auto_grow::yes),
       buffer_size(buffer.capacity()) // same rounding to BUFFER_WRAP
@@ -154,8 +146,8 @@ class SimpleWriter : public pyosmium::BaseHandler
     void close()
     {
         if (buffer) {
-            writer(std::move(buffer));
-            writer.close();
+            (*writer.get())(std::move(buffer));
+            writer.get()->close();
             buffer = osmium::memory::Buffer();
         }
     }
@@ -335,11 +327,11 @@ class SimpleWriter : public pyosmium::BaseHandler
             osmium::memory::Buffer new_buffer(buffer_size, osmium::memory::Buffer::auto_grow::yes);
             using std::swap;
             swap(buffer, new_buffer);
-            writer(std::move(new_buffer));
+            (*writer.get())(std::move(new_buffer));
         }
     }
 
-    osmium::io::Writer writer;
+    pyosmium::PyWriter writer;
     osmium::memory::Buffer buffer;
     size_t buffer_size;
 };
@@ -351,22 +343,38 @@ namespace pyosmium {
 void init_simple_writer(pybind11::module &m)
 {
     py::class_<SimpleWriter, BaseHandler>(m, "SimpleWriter")
-        .def(py::init<const char*, unsigned long, osmium::io::Header const *, bool, const std::string&>(),
+        .def(py::init<>([] (std::string file, unsigned long bufsz,
+                            osmium::io::Header const *header, bool overwrite,
+                            osmium::thread::Pool *pool) {
+                 return new SimpleWriter(osmium::io::File(std::move(file)), bufsz,
+                                         header, overwrite, pool);
+             }),
+             py::keep_alive<1, 6>(),
              py::arg("filename"), py::arg("bufsz") = 4096*1024,
              py::arg("header") = nullptr,
              py::arg("overwrite") = false,
-             py::arg("filetype") = "")
+             py::arg("thread_pool") = nullptr
+            )
         .def(py::init<>([] (std::filesystem::path const &file, unsigned long bufsz,
-                            osmium::io::Header const *header, bool overwrite) {
-                 return new SimpleWriter(file.string().c_str(), bufsz, header, overwrite, "");
+                            osmium::io::Header const *header, bool overwrite,
+                            osmium::thread::Pool *pool) {
+                 return new SimpleWriter(osmium::io::File(file.string()), bufsz,
+                                         header, overwrite, pool);
              }),
+             py::keep_alive<1, 6>(),
              py::arg("filename"), py::arg("bufsz") = 4096*1024,
              py::arg("header") = nullptr,
-             py::arg("overwrite") = false)
-        .def(py::init<osmium::io::File, unsigned long, osmium::io::Header const *, bool>(),
+             py::arg("overwrite") = false,
+             py::arg("thread_pool") = nullptr
+            )
+        .def(py::init<osmium::io::File, unsigned long, osmium::io::Header const *,
+                      bool, osmium::thread::Pool *>(),
+             py::keep_alive<1, 6>(),
              py::arg("filename"), py::arg("bufsz") = 4096*1024,
              py::arg("header") = nullptr,
-             py::arg("overwrite") = false)
+             py::arg("overwrite") = false,
+             py::arg("thread_pool") = nullptr
+            )
         .def("add_node", &SimpleWriter::add_node, py::arg("node"))
         .def("add_way", &SimpleWriter::add_way, py::arg("way"))
         .def("add_relation", &SimpleWriter::add_relation, py::arg("relation"))

src/osmium/_osmium.pyi

@@ -2,15 +2,15 @@
 #
 # This file is part of pyosmium. (https://osmcode.org/pyosmium/)
 #
-# Copyright (C) 2024 Sarah Hoffmann <[email protected]> and others.
+# Copyright (C) 2025 Sarah Hoffmann <[email protected]> and others.
 # For a full list of authors see the git log.
 from typing import ByteString, Union, Optional, Any
 import os
 
 from .osm import osm_entity_bits
 from .osm.types import OSMEntity
 from .index import LocationTable, IdSet
-from .io import Reader, Writer, Header, File, FileBuffer
+from .io import Reader, Writer, Header, File, FileBuffer, ThreadPool
 
 # Placeholder for more narrow type definition to come
 HandlerLike = object
@@ -127,18 +127,16 @@ class SimpleWriter(BaseHandler):
         when writing is finished.
     """
     def __init__(self, file: Union[str, 'os.PathLike[str]', File],
-                 bufsz: int= ...,
-                 header: Optional[Header]= ..., overwrite: bool= ...,
-                 filetype: str= ...) -> None:
+                 bufsz: int=4096*1024,
+                 header: Optional[Header]=None, overwrite: bool=False,
+                 thread_pool: Optional[ThreadPool]=None) -> None:
         """ Initiate a new writer for the given file. The writer will
             refuse to overwrite an already existing file unless _overwrite_
             is explicitly set to `True`.
 
             The file type is usually determined from the file extension.
             If you want to explicitly set the filetype (for example, when
             writing to standard output '-'), then use a File object.
-            Using the _filetype_ parameter to set the file type is deprecated
-            and only works when the file is a string.
 
             The _header_ parameter can be used to set a custom header in
             the output file. What kind of information can be written into
@@ -149,6 +147,14 @@ class SimpleWriter(BaseHandler):
             size is 4MB. Larger buffers are normally better but you should
             be aware that there are normally multiple buffers in use during
             the write process.
+
+            The writer implicitly creates a private
+            [ThreadPool][osmium.io.ThreadPool] which it may
+            use to parallelize writing to the output. Alternatively you
+            may hand in an externally created thread pool. This may be useful
+            when you create many writers in parallel and want them to share
+            a single thread pool or when you want to customize the size
+            of the thread pool.
         """
     def add_node(self, node: object) -> None:
         """ Add a new node to the file. The node may be a

src/osmium/back_reference_writer.py

@@ -2,15 +2,15 @@
 #
 # This file is part of pyosmium. (https://osmcode.org/pyosmium/)
 #
-# Copyright (C) 2024 Sarah Hoffmann <[email protected]> and others.
+# Copyright (C) 2025 Sarah Hoffmann <[email protected]> and others.
 # For a full list of authors see the git log.
-from typing import Any, Union
+from typing import Any, Union, Optional
 from pathlib import Path
 from tempfile import TemporaryDirectory
 import os
 
 from osmium._osmium import SimpleWriter
-from osmium.io import File, FileBuffer
+from osmium.io import File, FileBuffer, ThreadPool
 from osmium.file_processor import FileProcessor, zip_processors
 from osmium import IdTracker
 
@@ -30,7 +30,7 @@ class BackReferenceWriter:
     def __init__(self, outfile: Union[str, 'os.PathLike[str]', File],
                  ref_src: Union[str, 'os.PathLike[str]', File, FileBuffer],
                  overwrite: bool = False, remove_tags: bool = True,
-                 relation_depth: int = 0):
+                 relation_depth: int = 0, thread_pool: Optional[ThreadPool] = None):
         """ Create a new writer.
 
             `outfile` is the name of the output file to write. The file must
@@ -47,10 +47,17 @@ def __init__(self, outfile: Union[str, 'os.PathLike[str]', File],
             The writer will not complete nested relations by default. If you
             need nested relations, set `relation_depth` to the minimum depth
             to which relations shall be completed.
+
+            The writer implicitly creates a private
+            [ThreadPool][osmium.io.ThreadPool] which it
+            uses to parallelize IO operations. Alternatively you
+            may hand in an externally created thread pool.
         """
         self.outfile = outfile
         self.tmpdir = TemporaryDirectory()
-        self.writer = SimpleWriter(str(Path(self.tmpdir.name, 'back_writer.osm.pbf')))
+        self.thread_pool = thread_pool or ThreadPool()
+        self.writer = SimpleWriter(Path(self.tmpdir.name, 'back_writer.osm.pbf'),
+                                   thread_pool=self.thread_pool)
         self.overwrite = overwrite
         self.remove_tags = remove_tags
         self.id_tracker = IdTracker()
@@ -102,10 +109,13 @@ def close(self) -> None:
         self.id_tracker.complete_backward_references(self.ref_src,
                                                      relation_depth=self.relation_depth)
 
-        fp1 = FileProcessor(str(Path(self.tmpdir.name, 'back_writer.osm.pbf')))
-        fp2 = FileProcessor(self.ref_src).with_filter(self.id_tracker.id_filter())
+        fp1 = FileProcessor(Path(self.tmpdir.name, 'back_writer.osm.pbf'),
+                            thread_pool=self.thread_pool)
+        fp2 = FileProcessor(self.ref_src, thread_pool=self.thread_pool
+                            ).with_filter(self.id_tracker.id_filter())
 
-        with SimpleWriter(self.outfile, overwrite=self.overwrite) as writer:
+        with SimpleWriter(self.outfile, overwrite=self.overwrite,
+                          thread_pool=self.thread_pool) as writer:
             for o1, o2 in zip_processors(fp1, fp2):
                 if o1:
                     writer.add(o1)

src/osmium/file_processor.py

@@ -9,7 +9,7 @@
 
 import osmium
 from osmium.index import LocationTable
-from osmium.io import File, FileBuffer
+from osmium.io import File, FileBuffer, ThreadPool, Reader
 from osmium.osm.types import OSMEntity
 
 
@@ -20,7 +20,8 @@ class FileProcessor:
     """
 
     def __init__(self, indata: Union[File, FileBuffer, str, 'os.PathLike[str]'],
-                 entities: osmium.osm.osm_entity_bits = osmium.osm.ALL) -> None:
+                 entities: osmium.osm.osm_entity_bits = osmium.osm.ALL,
+                 thread_pool: Optional[ThreadPool] = None) -> None:
         """ Initialise a new file processor for the given input source _indata_.
             This may either be a filename, an instance of [File](IO.md#osmium.io.File)
             or buffered data in form of a [FileBuffer](IO.md#osmium.io.FileBuffer).
@@ -30,6 +31,13 @@ def __init__(self, indata: Union[File, FileBuffer, str, 'os.PathLike[str]'],
             directly at the source file and will never be passed to any filters
             including the location and area processors. You usually should not
             be restricting objects, when using those.
+
+            By default, pyosmium will create a private thread pool, which is
+            used to parallelize reading of the file. Alternatively you may
+            explicitly create a thread pool and hand it to the FileProcessor.
+            This may be necessary if you plan to run many processors in
+            parallel and want them to share a common thread pool or if you
+            want to change the default settings of the thread pool.
             """
         self._file = indata
         self._entities = entities
@@ -38,13 +46,14 @@ def __init__(self, indata: Union[File, FileBuffer, str, 'os.PathLike[str]'],
         self._filters: List['osmium._osmium.HandlerLike'] = []
         self._area_filters: List['osmium._osmium.HandlerLike'] = []
         self._filtered_handler: Optional['osmium._osmium.HandlerLike'] = None
+        self._thread_pool = thread_pool or ThreadPool()
 
     @property
     def header(self) -> osmium.io.Header:
         """ (read-only) [Header](IO.md#osmium.io.Header) information
             for the file to be read.
         """
-        return osmium.io.Reader(self._file, osmium.osm.NOTHING).header()
+        return Reader(self._file, osmium.osm.NOTHING, self._thread_pool).header()
 
     @property
     def node_location_storage(self) -> Optional[LocationTable]:
@@ -153,31 +162,28 @@ def __iter__(self) -> Iterator[OSMEntity]:
             handlers.append(lh)
 
         if self._area_handler is None:
-            reader = osmium.io.Reader(self._file, self._entities)
-            it = osmium.OsmFileIterator(reader, *handlers, *self._filters)
-            if self._filtered_handler:
-                it.set_filtered_handler(self._filtered_handler)
-            yield from it
+            with Reader(self._file, self._entities, thread_pool=self._thread_pool) as reader:
+                it = osmium.OsmFileIterator(reader, *handlers, *self._filters)
+                if self._filtered_handler:
+                    it.set_filtered_handler(self._filtered_handler)
+                yield from it
             return
 
         # need areas, do two pass handling
-        rd = osmium.io.Reader(self._file, osmium.osm.RELATION)
-        try:
+        with Reader(self._file, osmium.osm.RELATION, thread_pool=self._thread_pool) as rd:
             osmium.apply(rd, *self._area_filters, self._area_handler.first_pass_handler())
-        finally:
-            rd.close()
 
         buffer_it = osmium.BufferIterator(*self._filters)
         handlers.append(self._area_handler.second_pass_to_buffer(buffer_it))
 
-        reader = osmium.io.Reader(self._file, self._entities)
-        it = osmium.OsmFileIterator(reader, *handlers, *self._filters)
-        if self._filtered_handler:
-            it.set_filtered_handler(self._filtered_handler)
-        for obj in it:
-            yield obj
-            if buffer_it:
-                yield from buffer_it
+        with Reader(self._file, self._entities, thread_pool=self._thread_pool) as reader:
+            it = osmium.OsmFileIterator(reader, *handlers, *self._filters)
+            if self._filtered_handler:
+                it.set_filtered_handler(self._filtered_handler)
+            for obj in it:
+                yield obj
+                if buffer_it:
+                    yield from buffer_it
 
         # catch anything after the final flush
         if buffer_it:

src/osmium/forward_reference_writer.py

@@ -2,7 +2,7 @@
 #
 # This file is part of pyosmium. (https://osmcode.org/pyosmium/)
 #
-# Copyright (C) 2024 Sarah Hoffmann <[email protected]> and others.
+# Copyright (C) 2025 Sarah Hoffmann <[email protected]> and others.
 # For a full list of authors see the git log.
 from typing import Any, Optional, Union
 from pathlib import Path
@@ -11,7 +11,7 @@
 
 from osmium._osmium import SimpleWriter
 from osmium import IdTracker
-from osmium.io import File, FileBuffer
+from osmium.io import File, FileBuffer, ThreadPool
 from osmium.file_processor import FileProcessor, zip_processors
 
 
@@ -32,7 +32,8 @@ def __init__(self, outfile: Union[str, 'os.PathLike[str]', File],
                  ref_src: Union[str, 'os.PathLike[str]', File, FileBuffer],
                  overwrite: bool = False, back_references: bool = True,
                  remove_tags: bool = True, forward_relation_depth: int = 0,
-                 backward_relation_depth: int = 1) -> None:
+                 backward_relation_depth: int = 1,
+                 thread_pool: Optional[ThreadPool] = None) -> None:
         """ Create a new writer.
 
             `outfile` is the name of the output file to write. The file must
@@ -52,7 +53,9 @@ def __init__(self, outfile: Union[str, 'os.PathLike[str]', File],
         """
         self.outfile = outfile
         self.tmpdir: Optional['TemporaryDirectory[Any]'] = TemporaryDirectory()
-        self.writer = SimpleWriter(str(Path(self.tmpdir.name, 'forward_writer.osm.pbf')))
+        self.thread_pool = thread_pool or ThreadPool()
+        self.writer = SimpleWriter(Path(self.tmpdir.name, 'forward_writer.osm.pbf'),
+                                   thread_pool=self.thread_pool)
         self.overwrite = overwrite
         self.back_references = back_references
         self.id_tracker = IdTracker()
@@ -118,10 +121,13 @@ def close(self) -> None:
                     self.ref_src,
                     relation_depth=self.backward_relation_depth)
 
-            fp1 = FileProcessor(Path(self.tmpdir.name, 'forward_writer.osm.pbf'))
-            fp2 = FileProcessor(self.ref_src).with_filter(self.id_tracker.id_filter())
+            fp1 = FileProcessor(Path(self.tmpdir.name, 'forward_writer.osm.pbf'),
+                                thread_pool=self.thread_pool)
+            fp2 = FileProcessor(self.ref_src, thread_pool=self.thread_pool
+                                ).with_filter(self.id_tracker.id_filter())
 
-            with SimpleWriter(self.outfile, overwrite=self.overwrite) as writer:
+            with SimpleWriter(self.outfile, overwrite=self.overwrite,
+                              thread_pool=self.thread_pool) as writer:
                 for o1, o2 in zip_processors(fp1, fp2):
                     if o1:
                         writer.add(o1)

src/osmium/helper.py

@@ -56,8 +56,8 @@ class WriteHandler(SimpleWriter):
         documentation.
     """
 
-    def __init__(self, filename: str, bufsz: int = 4096*1024, filetype: str = "") -> None:
-        super().__init__(filename, bufsz=bufsz, filetype=filetype)
+    def __init__(self, filename: str, bufsz: int = 4096*1024) -> None:
+        super().__init__(filename, bufsz=bufsz)
 
 
 def _merge_apply(self: MergeInputReader, *handlers: 'HandlerLike',