update typing information

Also fixes a number of typing issues mypy found.

src/osmium/__init__.py

@@ -27,11 +27,11 @@ class WriteHandler(SimpleWriter):
         documentation.
     """
 
-    def __init__(filename: str, bufsz: int=4096*1024, filetype: str="") -> None:
+    def __init__(self, filename: str, bufsz: int=4096*1024, filetype: str="") -> None:
         super().__init__(filename, bufsz=bufsz, filetype=filetype)
 
 
-def _merge_apply(self, *handlers: Any, idx: str = '', simplify: bool = True) -> None:
+def _merge_apply(self: osmium.MergeInputReader, *handlers: Any, idx: str = '', simplify: bool = True) -> None:
     """ Apply collected data to a handler. The data will be sorted first.
         If `simplify` is true (default) then duplicates will be eliminated
         and only the newest version of each object kept. If `idx` is given
@@ -51,4 +51,4 @@ def _merge_apply(self, *handlers: Any, idx: str = '', simplify: bool = True) ->
 
     self._apply_internal(*handlers, simplify=simplify)
 
-osmium.MergeInputReader.apply = _merge_apply
+osmium.MergeInputReader.apply = _merge_apply # type: ignore[method-assign]

src/osmium/_osmium.pyi

@@ -1,55 +1,94 @@
-from typing import overload, ByteString, Union
+# SPDX-License-Identifier: BSD-2-Clause
+#
+# This file is part of pyosmium. (https://osmcode.org/pyosmium/)
+#
+# Copyright (C) 2024 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
 
-import osmium.index
-import osmium.io
+from .osm import osm_entity_bits
+from .osm.types import OSMEntity
+from .index import LocationTable, IdSet
+from .io import Reader, Writer, Header
 
 StrPath = Union[str, 'os.PathLike[str]']
 
+# Placeholder for more narrow type defintion to come
+HandlerLike = object
+
 class InvalidLocationError(Exception): ...
 
-class NodeLocationsForWays:
-    def __init__(self, locations: osmium.index.LocationTable) -> None: ...
-    def ignore_errors(self) -> None: ...
 
 class BaseHandler: ...
 
-class SimpleHandler(BaseHandler):
-    def __init__(self) -> None: ...
-    def apply_buffer(self, buffer: Union[ByteString, str], format: str, locations: bool = ..., idx: str = ...) -> None: ...
-    def apply_file(self, filename: StrPath, locations: bool = ..., idx: str = ...) -> None: ...
+
+class BaseFilter(BaseHandler):
+    def enable_for(self, entities: osm_entity_bits) -> None: ...
+
+
+class BufferIterator:
+    def __init__(self, *handlers: HandlerLike) -> None: ...
+    def __bool__(self) -> bool: ...
+    def __iter__(self) -> 'BufferIterator': ...
+    def __next__(self) -> OSMEntity: ...
+
 
 class MergeInputReader:
     def __init__(self) -> None: ...
+    def _apply_internal(self, *handlers: HandlerLike, simplify: bool = ...) -> None: ...
     def add_buffer(self, buffer: Union[ByteString, str], format: str) -> int: ...
     def add_file(self, file: str) -> int: ...
-    def apply(self, handler: BaseHandler, idx: str = ..., simplify: bool = ...) -> None: ...
-    def apply_to_reader(self, reader: osmium.io.Reader, writer: osmium.io.Writer, with_history: bool = ...) -> None: ...
-
-class WriteHandler(BaseHandler):
-    @overload
-    def __init__(self, filename: str, bufsz: int, filetype: str) -> None: ...
-    @overload
-    def __init__(self, filename: str, bufsz: int) -> None: ...
-    @overload
-    def __init__(self, filename: str) -> None: ...
-    def close(self) -> None: ...
+    def apply_to_reader(self, reader: Reader, writer: Writer, with_history: bool = ...) -> None: ...
+    def apply(self, *handlers: Any, idx: str = '', simplify: bool = True) -> None: ...
+
+
 
 class SimpleWriter:
-    @overload
-    def __init__(self, filename: str, bufsz: int, header: osmium.io.Header) -> None: ...
-    @overload
-    def __init__(self, filename: str, bufsz: int) -> None: ...
-    @overload
-    def __init__(self, filename: str) -> None: ...
+    def __init__(self, filename: str, bufsz: int= ...,
+                 header: Optional[Header]= ..., overwrite: bool= ...,
+                 filetype: str= ...) -> None: ...
     def add_node(self, node: object) -> None: ...
     def add_relation(self, relation: object) -> None: ...
     def add_way(self, way: object) -> None: ...
+    def add(self, obj: object) -> None: ...
     def close(self) -> None: ...
+    def __enter__(self) -> 'SimpleWriter':...
+    def __exit__(self, *args: Any) -> None:...
+
+
+class NodeLocationsForWays:
+    apply_nodes_to_ways: bool
+    def __init__(self, locations: LocationTable) -> None: ...
+    def ignore_errors(self) -> None: ...
+
+
+class OsmFileIterator:
+    def __init__(self, reader: Reader, *handlers: HandlerLike) -> None: ...
+    def set_filtered_handler(self, handler: object) -> None: ...
+    def __iter__(self) -> 'OsmFileIterator': ...
+    def __next__(self) -> OSMEntity: ...
+
+
+class IdTrackerIdFilter(BaseFilter): ...
+
+
+class IdTrackerContainsFilter(BaseFilter): ...
+
+
+class IdTracker:
+    def __init__(self) -> None: ...
+    def add_node(self, node: int) -> None: ...
+    def add_relation(self, relation: int) -> None: ...
+    def add_way(self, way: int) -> None: ...
+    def add_references(self, obj: object) -> None: ...
+    def contains_any_references(self, obj: object) -> bool: ...
+    def complete_backward_references(self, filename: str, relation_depth: int = ...) -> None: ...
+    def complete_forward_references(self, filename: str, relation_depth: int = ...) -> None: ...
+    def id_filter(self) -> IdTrackerIdFilter: ...
+    def contains_filter(self) -> IdTrackerContainsFilter: ...
+    def node_ids(self) -> IdSet: ...
+    def way_ids(self) -> IdSet: ...
+    def relation_ids(self) -> IdSet: ...
 
-@overload
-def apply(reader: osmium.io.Reader, handler: BaseHandler) -> None: ...
-@overload
-def apply(reader: osmium.io.Reader, node_handler: NodeLocationsForWays) -> None: ...
-@overload
-def apply(reader: osmium.io.Reader, node_handler: NodeLocationsForWays, handler: BaseHandler) -> None: ...
+def apply(reader: Union[Reader | str], *handlers: HandlerLike) -> None: ...

src/osmium/area.pyi

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: BSD-2-Clause
+#
+# This file is part of pyosmium. (https://osmcode.org/pyosmium/)
+#
+# Copyright (C) 2024 Sarah Hoffmann <[email protected]> and others.
+# For a full list of authors see the git log.
+
+from . import BaseHandler, HandlerLike, BufferIterator
+
+class AreaManagerSecondPassHandler(BaseHandler): ...
+
+class AreaManagerBufferHandler(BaseHandler):...
+
+
+class AreaManager(BaseHandler):
+    def __init__(self) -> None:...
+    def first_pass_handler(self) -> 'AreaManager':...
+    def second_pass_handler(self, *handlers: HandlerLike) -> AreaManagerSecondPassHandler:...
+    def second_pass_to_buffer(self, callback: BufferIterator) -> AreaManagerBufferHandler:...

src/osmium/back_reference_writer.py

@@ -103,7 +103,10 @@ def close(self) -> None:
             for o1, o2 in zip_processors(fp1, fp2):
                 if o1:
                     writer.add(o1)
-                else:
-                    writer.add(o2.replace(tags={}) if self.remove_tags else o2)
+                elif o2:
+                    if self.remove_tags and hasattr(o2, 'replace'):
+                        writer.add(o2.replace(tags={}))
+                    else:
+                        writer.add(o2)
 
         self.tmpdir.cleanup()

src/osmium/file_processor.py

@@ -4,43 +4,46 @@
 #
 # Copyright (C) 2024 Sarah Hoffmann <[email protected]> and others.
 # For a full list of authors see the git log.
-from typing import Iterable, Tuple, Any
+from typing import Iterable, Iterator, Tuple, Any, Union, Optional, List
 from pathlib import Path
 
 import osmium
+from osmium.index import LocationTable
+from osmium.osm.types import OSMEntity
 
 class FileProcessor:
     """ A generator that emits OSM objects read from a file.
     """
 
-    def __init__(self, filename, entities=osmium.osm.ALL):
+    def __init__(self, filename: Union[osmium.io.File, osmium.io.FileBuffer, str, Path],
+                 entities: osmium.osm.osm_entity_bits=osmium.osm.ALL) -> None:
         if isinstance(filename, (osmium.io.File, osmium.io.FileBuffer)):
             self._file = filename
         elif isinstance(filename, (str, Path)):
             self._file = osmium.io.File(str(filename))
         else:
             raise TypeError("File must be an osmium.io.File, osmium.io.FileBuffer, str or Path")
         self._entities = entities
-        self._node_store = None
-        self._area_handler = None
-        self._filters = []
-        self._area_filters = []
-        self._filtered_handler = None
+        self._node_store: Optional[LocationTable] = None
+        self._area_handler: Optional[osmium.area.AreaManager] = None
+        self._filters: List['osmium._osmium.HandlerLike'] = []
+        self._area_filters: List['osmium._osmium.HandlerLike'] = []
+        self._filtered_handler: Optional['osmium._osmium.HandlerLike'] = None
 
     @property
-    def header(self):
+    def header(self) -> osmium.io.Header:
         """ Return the header information for the file to be read.
         """
         return osmium.io.Reader(self._file, osmium.osm.NOTHING).header()
 
     @property
-    def node_location_storage(self):
+    def node_location_storage(self) -> Optional[LocationTable]:
         """ Return the node location cache, if enabled.
             This can be used to manually look up locations of nodes.
         """
         return self._node_store
 
-    def with_locations(self, storage='flex_mem'):
+    def with_locations(self, storage: str='flex_mem') -> 'FileProcessor':
         """ Enable caching of node locations. This is necessary in order
             to get geometries for ways and relations.
         """
@@ -55,7 +58,7 @@ def with_locations(self, storage='flex_mem'):
 
         return self
 
-    def with_areas(self, *filters):
+    def with_areas(self, *filters: 'osmium._osmium.HandlerLike') -> 'FileProcessor':
         """ Enable area processing. When enabled, then closed ways and
             relations of type multipolygon will also be returned as an
             Area type.
@@ -81,7 +84,7 @@ def with_areas(self, *filters):
         self._area_filters.extend(filters)
         return self
 
-    def with_filter(self, filt):
+    def with_filter(self, filt: 'osmium._osmium.HandlerLike') -> 'FileProcessor':
         """ Add a filter function that is called before an object is
             returned in the iterator. Filters are applied sequentially
             in the order they were added.
@@ -90,17 +93,17 @@ def with_filter(self, filt):
         return self
 
 
-    def handler_for_filtered(self, handler):
+    def handler_for_filtered(self, handler: 'osmium._osmium.HandlerLike') -> 'FileProcessor':
         """ Set a handler to be called on all objects that have been
             filtered out and are not presented to the iterator loop.
         """
         self._filtered_handler = handler
         return self
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[OSMEntity]:
         """ Return the iterator over the file.
         """
-        handlers = []
+        handlers: List['osmium._osmium.HandlerLike'] = []
 
         if self._node_store is not None:
             lh = osmium.NodeLocationsForWays(self._node_store)
@@ -139,7 +142,7 @@ def __iter__(self):
             yield from buffer_it
 
 
-def zip_processors(*procs: FileProcessor) -> Iterable[Tuple[Any, ...]]:
+def zip_processors(*procs: FileProcessor) -> Iterable[List[Optional[OSMEntity]]]:
     """ Return the data from the FileProcessors in parallel such
         that objects with the same ID are returned at the same time.
 
@@ -150,19 +153,19 @@ def zip_processors(*procs: FileProcessor) -> Iterable[Tuple[Any, ...]]:
 
     class _CompIter:
 
-        def __init__(self, fp):
+        def __init__(self, fp: FileProcessor) -> None:
             self.iter = iter(fp)
-            self.current = None
-            self.comp = None
+            self.current: Optional[OSMEntity] = None
+            self.comp: Optional[Tuple[int, int]] = None
 
-        def val(self, nextid):
+        def val(self, nextid: Tuple[int, int]) -> Optional[OSMEntity]:
             """ Return current object if it corresponds to the given object ID.
             """
             if self.comp == nextid:
                 return self.current
             return None
 
-        def next(self, nextid):
+        def next(self, nextid: Optional[Tuple[int, int]]) -> Tuple[int, int]:
             """ Get the next object ID, if and only if nextid points to the
                 previously returned object ID. Otherwise return the previous
                 ID again.
@@ -173,6 +176,7 @@ def next(self, nextid):
                     self.comp = (100, 0) # end of file marker. larger than any ID
                 else:
                     self.comp = (TID[self.current.type_str()], self.current.id)
+            assert self.comp is not None
             return self.comp
 
 
@@ -181,6 +185,5 @@ def next(self, nextid):
     nextid = min(i.next(None) for i in iters)
 
     while nextid[0] < 100:
-        yield (i.val(nextid) for i in iters)
-
+        yield [i.val(nextid) for i in iters]
         nextid = min(i.next(nextid) for i in iters)

src/osmium/filter.pyi

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: BSD-2-Clause
+#
+# This file is part of pyosmium. (https://osmcode.org/pyosmium/)
+#
+# Copyright (C) 2024 Sarah Hoffmann <[email protected]> and others.
+# For a full list of authors see the git log.
+from typing import Tuple, Iterable
+
+from ._osmium import BaseFilter
+from .osm import osm_entity_bits
+
+class EmptyTagFilter(BaseFilter):
+    def __init__(self) -> None: ...
+
+
+class KeyFilter(BaseFilter):
+    def __init__(self, *keys: str) -> None: ...
+
+
+class TagFilter(BaseFilter):
+    def __init__(self, *tags: Tuple[str, str]) -> None: ...
+
+
+class EntityFilter(BaseFilter):
+    def __init__(self, entities: osm_entity_bits) -> None: ...
+
+
+class IdFilter(BaseFilter):
+    def __init__(self, ids: Iterable[int]) -> None: ...
+
+
+class GeoInterfaceFilter(BaseFilter):
+    def __init__(self, drop_invalid_geometries: bool= ..., tags: Iterable[str] = ...) -> None: ...