File: file_iterator.cc

package info (click to toggle)
pyosmium 4.2.0-2
  • links: PTS, VCS
  • area: main
  • in suites: forky
  • size: 5,584 kB
  • sloc: python: 4,400; cpp: 2,504; makefile: 20
file content (154 lines) | stat: -rw-r--r-- 4,987 bytes parent folder | download | duplicates (5)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
/* SPDX-License-Identifier: BSD-2-Clause
 *
 * This file is part of pyosmium. (https://osmcode.org/pyosmium/)
 *
 * Copyright (C) 2024 Sarah Hoffmann <lonvia@denofr.de> and others.
 * For a full list of authors see the git log.
 */
#include <pybind11/pybind11.h>

#include <osmium/io/any_input.hpp>
#include <osmium/io/file.hpp>

#include "osmium_module.h"
#include "base_handler.h"
#include "osm_base_objects.h"
#include "handler_chain.h"
#include "python_handler.h"

namespace py = pybind11;

namespace {

class OsmFileIterator
{
public:
    OsmFileIterator(osmium::io::Reader *reader, py::args args)
    : m_reader(reader), m_handler(args)
    {
        m_buffer = m_reader->read();

        if (m_buffer) {
            m_buffer_it = m_buffer.begin();
        }
    }

    pybind11::object next()
    {
        while (true) {
            m_current.emplace<bool>(false);

            if (!m_buffer) {
                throw pybind11::stop_iteration();
            }

            while (m_buffer_it == m_buffer.end()) {
                m_buffer = m_reader->read();
                if (!m_buffer) {
                    m_handler.flush();
                    throw pybind11::stop_iteration();
                }
                m_buffer_it = m_buffer.begin();
            }

            osmium::OSMEntity *entity = &*m_buffer_it;
            ++m_buffer_it;

            switch (entity->type()) {
                case osmium::item_type::node:
                {
                    auto &obj = m_current.emplace<pyosmium::PyOSMNode>(entity);
                    if (!m_handler.node(obj)) {
                        return obj.get_or_create_python_object();
                    } else if (m_filtered_handler) {
                        m_filtered_handler->node(obj);
                    }
                    break;
                }
                case osmium::item_type::way:
                {
                    auto &obj = m_current.emplace<pyosmium::PyOSMWay>(entity);
                    if (!m_handler.way(obj)) {
                        return obj.get_or_create_python_object();
                    } else if (m_filtered_handler) {
                        m_filtered_handler->way(obj);
                    }
                    break;
                }
                case osmium::item_type::relation:
                {
                    auto &obj = m_current.emplace<pyosmium::PyOSMRelation>(entity);
                    if (!m_handler.relation(obj)) {
                        return obj.get_or_create_python_object();
                    } else if (m_filtered_handler) {
                        m_filtered_handler->relation(obj);
                    }
                    break;
                }
                case osmium::item_type::area:
                {
                    auto &obj = m_current.emplace<pyosmium::PyOSMArea>(entity);
                    if (!m_handler.area(obj)) {
                        return obj.get_or_create_python_object();
                    } else if (m_filtered_handler) {
                        m_filtered_handler->area(obj);
                    }
                    break;
                }
                case osmium::item_type::changeset:
                {
                    auto &obj = m_current.emplace<pyosmium::PyOSMChangeset>(entity);
                    if (!m_handler.changeset(obj)) {
                        return obj.get_or_create_python_object();
                    } else if (m_filtered_handler) {
                        m_filtered_handler->changeset(obj);
                    }
                    break;
                }
                default:
                    break;
            }
       }

       return pybind11::object();
    }

    void set_filtered_handler(pyosmium::BaseHandler *handler) {
        m_filtered_handler = handler;
    }

    void set_filtered_python_handler(pybind11::handle handler) {
        m_filtered_python_handler = std::make_unique<pyosmium::PythonHandler>(handler);
        m_filtered_handler = m_filtered_python_handler.get();
    }

private:
    osmium::io::Reader *m_reader;
    osmium::memory::Buffer m_buffer;
    osmium::memory::Buffer::iterator m_buffer_it;
    pyosmium::PyOSMAny m_current;

    pyosmium::HandlerChain m_handler;
    pyosmium::BaseHandler *m_filtered_handler = nullptr;
    std::unique_ptr<pyosmium::PythonHandler> m_filtered_python_handler;
};

} // namespace

namespace pyosmium {

void init_osm_file_iterator(py::module &m)
{
    py::class_<OsmFileIterator>(m, "OsmFileIterator")
        .def(py::init<osmium::io::Reader *, py::args>(),
             py::keep_alive<0, 1>())
        .def("set_filtered_handler", &OsmFileIterator::set_filtered_handler,
             py::keep_alive<0, 1>())
        .def("set_filtered_handler", &OsmFileIterator::set_filtered_python_handler,
             py::keep_alive<0, 1>())
        .def("__iter__", [](py::object const &self) { return self; })
        .def("__next__", &OsmFileIterator::next)
        ;
}

} // namespace